import requests
import time
import os
import sys
from sqlalchemy import create_engine, ForeignKey, UniqueConstraint, Index
from sqlalchemy.orm import sessionmaker, relationship, DeclarativeBase, mapped_column, Mapped
from dataclasses import dataclass
from typing import List

class Base(DeclarativeBase):
    pass

@dataclass
class Host(Base):
    __tablename__ = 'hosts'
    id: Mapped[int] = mapped_column(primary_key=True)
    hostname: Mapped[str] = mapped_column(unique=True)
    current_stats: Mapped[List["CurrentStat"]] = relationship("CurrentStat", back_populates="host")
    historical_stats: Mapped[List["HistoricalStat"]] = relationship("HistoricalStat", back_populates="host")

@dataclass
class CurrentStat(Base):
    __tablename__ = 'current_stats'
    id: Mapped[int] = mapped_column(primary_key=True)
    host_id: Mapped[int] = mapped_column(ForeignKey('hosts.id'))
    category: Mapped[str] = mapped_column()
    name: Mapped[str] = mapped_column()
    value: Mapped[str] = mapped_column()
    unit: Mapped[str | None] = mapped_column(nullable=True)
    timestamp: Mapped[float] = mapped_column()
    host: Mapped["Host"] = relationship("Host", back_populates="current_stats")
    __table_args__ = (
        UniqueConstraint('host_id', 'category', 'name'),
        Index('idx_current_host', 'host_id'),
    )

@dataclass
class HistoricalStat(Base):
    __tablename__ = 'historical_stats'
    id: Mapped[int] = mapped_column(primary_key=True)
    host_id: Mapped[int] = mapped_column(ForeignKey('hosts.id'))
    category: Mapped[str] = mapped_column()
    name: Mapped[str] = mapped_column()
    value: Mapped[str] = mapped_column()
    unit: Mapped[str | None] = mapped_column(nullable=True)
    timestamp: Mapped[float] = mapped_column()
    host: Mapped["Host"] = relationship("Host", back_populates="historical_stats")
    __table_args__ = (
        Index('idx_historical_host', 'host_id'),
        Index('idx_historical_timestamp', 'timestamp'),
    )

class StatsCollector:
    def __init__(self, db_path="stats.db"):
        state_home = os.environ.get("XDG_STATE_HOME", os.path.expanduser("~/.local/state"))
        stats_dir = os.path.join(state_home, "stats")
        os.makedirs(stats_dir, exist_ok=True)
        self.db_path = os.path.join(stats_dir, db_path)
        if os.path.exists(self.db_path):
            os.remove(self.db_path)
        self.engine = create_engine(f'sqlite:///{self.db_path}')
        self.setup_database()
        Session = sessionmaker(bind=self.engine)
        self.session = Session()

    def setup_database(self):
        """Initialize the SQLite database schema"""
        Base.metadata.create_all(self.engine)

    def get_or_create_host_id(self, hostname):
        """Get host ID from database or create if not exists"""
        host = self.session.query(Host).filter_by(hostname=hostname).first()
        if not host:
            host = Host(hostname=hostname)
            self.session.add(host)
            self.session.commit()
        return host.id

    def store_stats(self, hostname, stats_data):
        """Store stats in the database"""
        host_id = self.get_or_create_host_id(hostname)
        timestamp = stats_data.get('timestamp', time.time())

        for category, category_stats in stats_data['stats'].items():
            for name, data in category_stats.items():
                value = data['value']
                unit = data.get('unit', None)
                is_historical = data.get('track', False)

                if is_historical:
                    hist_stat = HistoricalStat(
                        host_id=host_id,
                        category=category,
                        name=name,
                        value=value,
                        unit=unit,
                        timestamp=timestamp
                    )
                    self.session.add(hist_stat)
                else:
                    # Only update/create current stat if not historical
                    curr_stat = self.session.query(CurrentStat).filter_by(
                        host_id=host_id,
                        category=category,
                        name=name
                    ).first()

                    if curr_stat:
                        curr_stat.value = value
                        curr_stat.unit = unit
                        curr_stat.timestamp = timestamp
                    else:
                        curr_stat = CurrentStat(
                            host_id=host_id,
                            category=category,
                            name=name,
                            value=value,
                            unit=unit,
                            timestamp=timestamp
                        )
                        self.session.add(curr_stat)

        self.session.commit()

    def collect_from_host(self, hostname, port=5747):
        """Collect stats from a single host"""
        try:
            url = f"http://{hostname}:{port}"
            response = requests.get(url, timeout=5)
            response.raise_for_status()
            stats_data = response.json()
            self.store_stats(hostname, stats_data)
            print(f"Successfully collected stats from {hostname}")
            return True
        except Exception as e:
            print(f"Error collecting stats from {hostname}: {str(e)}", file=sys.stderr)
            return False