"""Cost tracking service for per-session and per-plan LLM spending."""

from __future__ import annotations

from datetime import UTC, datetime
from pathlib import Path

from pydantic import BaseModel
from sqlalchemy import DateTime, Float, Integer, String, create_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker


class _Base(DeclarativeBase):
    pass


class CostRecord(BaseModel):
    """Record of LLM API usage and cost."""

    session_id: str
    plan_id: str
    tokens_in: int
    tokens_out: int
    model: str
    cost_usd: float
    timestamp: datetime


class CostEntry(_Base):
    """SQLAlchemy model for cost tracking."""

    __tablename__ = "cost_entries"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    session_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    plan_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    tokens_in: Mapped[int] = mapped_column(Integer, nullable=False)
    tokens_out: Mapped[int] = mapped_column(Integer, nullable=False)
    model: Mapped[str] = mapped_column(String, nullable=False)
    cost_usd: Mapped[float] = mapped_column(Float, nullable=False)
    timestamp: Mapped[datetime] = mapped_column(DateTime, nullable=False)


class CostTracker:
    """Service for tracking LLM API spending per session and per plan."""

    def __init__(
        self,
        data_dir: Path,
        model_pricing: dict[str, dict[str, float]] | None = None,
    ) -> None:
        """Initialize the CostTracker.

        Args:
            data_dir: Directory where costs.db will be stored
            model_pricing: Dict mapping model names to pricing dicts with
                          'input' and 'output' keys (cost per 1K tokens)
        """
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.db_path = self.data_dir / "costs.db"

        self.model_pricing = model_pricing or {
            "gpt-4": {"input": 0.03, "output": 0.06},
            "gpt-4-turbo": {"input": 0.01, "output": 0.03},
            "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
            "claude-3-opus": {"input": 0.015, "output": 0.075},
            "claude-3-sonnet": {"input": 0.003, "output": 0.015},
            "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
        }

        self._init_db()

    def _init_db(self) -> None:
        """Initialize the SQLite database."""
        engine = create_engine(f"sqlite:///{self.db_path}")
        _Base.metadata.create_all(engine)
        self.engine = engine
        self.SessionLocal = sessionmaker(bind=engine)

    def _calculate_cost(self, tokens_in: int, tokens_out: int, model: str) -> float:
        """Calculate cost for given tokens and model.

        Args:
            tokens_in: Number of input tokens
            tokens_out: Number of output tokens
            model: Model name

        Returns:
            Cost in USD
        """
        pricing = self.model_pricing.get(model, {"input": 0.0, "output": 0.0})
        input_cost = (tokens_in / 1000) * pricing.get("input", 0.0)
        output_cost = (tokens_out / 1000) * pricing.get("output", 0.0)
        return round(input_cost + output_cost, 6)

    def record_usage(
        self,
        session_id: str,
        plan_id: str,
        tokens_in: int,
        tokens_out: int,
        model: str,
    ) -> CostRecord:
        """Record LLM API usage.

        Args:
            session_id: Session identifier
            plan_id: Plan identifier
            tokens_in: Number of input tokens
            tokens_out: Number of output tokens
            model: Model name

        Returns:
            CostRecord with calculated cost
        """
        cost_usd = self._calculate_cost(tokens_in, tokens_out, model)
        timestamp = datetime.now(UTC).replace(tzinfo=None)

        session = self.SessionLocal()
        try:
            entry = CostEntry(
                session_id=session_id,
                plan_id=plan_id,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
                model=model,
                cost_usd=cost_usd,
                timestamp=timestamp,
            )
            session.add(entry)
            session.commit()
        finally:
            session.close()

        return CostRecord(
            session_id=session_id,
            plan_id=plan_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            model=model,
            cost_usd=cost_usd,
            timestamp=timestamp,
        )

    def get_session_cost(self, session_id: str) -> float:
        """Get total cost for a session.

        Args:
            session_id: Session identifier

        Returns:
            Total cost in USD for the session
        """
        session = self.SessionLocal()
        try:
            entries = (
                session.query(CostEntry)
                .filter(CostEntry.session_id == session_id)
                .all()
            )
            return round(sum(entry.cost_usd for entry in entries), 6)
        finally:
            session.close()

    def get_plan_cost(self, plan_id: str) -> float:
        """Get total cost for a plan.

        Args:
            plan_id: Plan identifier

        Returns:
            Total cost in USD for the plan
        """
        session = self.SessionLocal()
        try:
            entries = (
                session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all()
            )
            return round(sum(entry.cost_usd for entry in entries), 6)
        finally:
            session.close()

    def get_session_entries(self, session_id: str) -> list[CostRecord]:
        """Get all cost entries for a session.

        Args:
            session_id: Session identifier

        Returns:
            List of CostRecord objects for the session
        """
        session = self.SessionLocal()
        try:
            entries = (
                session.query(CostEntry)
                .filter(CostEntry.session_id == session_id)
                .all()
            )
            return [
                CostRecord(
                    session_id=entry.session_id,
                    plan_id=entry.plan_id,
                    tokens_in=entry.tokens_in,
                    tokens_out=entry.tokens_out,
                    model=entry.model,
                    cost_usd=entry.cost_usd,
                    timestamp=entry.timestamp,
                )
                for entry in entries
            ]
        finally:
            session.close()

    def get_plan_entries(self, plan_id: str) -> list[CostRecord]:
        """Get all cost entries for a plan.

        Args:
            plan_id: Plan identifier

        Returns:
            List of CostRecord objects for the plan
        """
        session = self.SessionLocal()
        try:
            entries = (
                session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all()
            )
            return [
                CostRecord(
                    session_id=entry.session_id,
                    plan_id=entry.plan_id,
                    tokens_in=entry.tokens_in,
                    tokens_out=entry.tokens_out,
                    model=entry.model,
                    cost_usd=entry.cost_usd,
                    timestamp=entry.timestamp,
                )
                for entry in entries
            ]
        finally:
            session.close()
