"""Verbrauch erfassen und auswerten (→ Konzept 8.1, 13)."""

import uuid
from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
from typing import Any

from sqlalchemy import case, func, select
from sqlalchemy.orm import Session

from app.ai.models import Model
from app.identity.models import Team, User
from app.tasks.models import TaskRun
from app.usage.models import UsageRecord
from app.workspaces.models import Workspace

MILLION = Decimal(1_000_000)
GROUPS = {"team": "Team", "workspace": "Workspace", "model": "Modell", "user": "Nutzer", "day": "Tag"}


def tokens(usage: dict[str, Any]) -> dict[str, int]:
    """Responses-API-Format: input_tokens, output_tokens und die Details zu Cache und Reasoning."""
    def number(value: Any) -> int:
        return value if isinstance(value, int) and value >= 0 else 0
    return {"input_tokens": number(usage.get("input_tokens")),
            "cached_tokens": number((usage.get("input_tokens_details") or {}).get("cached_tokens")),
            "output_tokens": number(usage.get("output_tokens")),
            "reasoning_tokens": number((usage.get("output_tokens_details") or {}).get("reasoning_tokens"))}


def price(model: Model | None, counts: dict[str, int]) -> Decimal | None:
    """Gecachte Eingabe-Tokens zum Cache-Preis (falls hinterlegt), Reasoning zählt zur Ausgabe."""
    if model is None or model.input_price is None or model.output_price is None:
        return None
    cached = min(counts["cached_tokens"], counts["input_tokens"])
    cached_price = model.cached_price if model.cached_price is not None else model.input_price
    return ((counts["input_tokens"] - cached) * model.input_price + cached * cached_price
            + counts["output_tokens"] * model.output_price) / MILLION


def record(session: Session, *, run_id: uuid.UUID | None, workspace_id: int | None, user_id: int | None,
           model_id: int | None, model_name: str, usage: dict[str, Any]) -> UsageRecord:
    counts = tokens(usage)
    workspace = session.get(Workspace, workspace_id) if workspace_id else None
    model = session.get(Model, model_id) if model_id else None
    entry = UsageRecord(run_id=run_id, workspace_id=workspace_id, team_id=workspace.team_id if workspace else None,
                        user_id=user_id, model=model_name, cost=price(model, counts), **counts)
    session.add(entry)
    return entry


@dataclass
class Row:
    key: Any
    label: str
    calls: int
    runs: int
    input_tokens: int
    cached_tokens: int
    output_tokens: int
    reasoning_tokens: int
    cost: Decimal | None
    unpriced: int  # Aufrufe ohne Preis

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


def report(session: Session, start: datetime, end: datetime, group: str, *, team_ids: set[int] | None = None,
           workspace_id: int | None = None) -> list[Row]:
    """Verbrauch im Zeitraum [start, end), gruppiert; team_ids/workspace_id schränken ein."""
    if group not in GROUPS:
        raise ValueError(group)
    column = {"team": UsageRecord.team_id, "workspace": UsageRecord.workspace_id, "model": UsageRecord.model,
              "user": UsageRecord.user_id,
              "day": func.date_trunc("day", UsageRecord.time)}[group].label("key")
    query = (select(column, func.count(), func.count(func.distinct(UsageRecord.run_id)),
                    func.sum(UsageRecord.input_tokens), func.sum(UsageRecord.cached_tokens),
                    func.sum(UsageRecord.output_tokens), func.sum(UsageRecord.reasoning_tokens),
                    func.sum(UsageRecord.cost), func.sum(case((UsageRecord.cost.is_(None), 1), else_=0)))
             .where(UsageRecord.time >= start, UsageRecord.time < end).group_by(column))
    if team_ids is not None:
        query = query.where(UsageRecord.team_id.in_(team_ids))
    if workspace_id is not None:
        query = query.where(UsageRecord.workspace_id == workspace_id)
    rows = session.execute(query).all()
    labels = names(session, group, {r[0] for r in rows})
    result = [Row(key, labels.get(key, str(key)), calls, runs, int(i or 0), int(c or 0), int(o or 0), int(r or 0),
                  cost, int(unpriced or 0))
              for key, calls, runs, i, c, o, r, cost, unpriced in rows]
    if group == "day":
        return sorted(result, key=lambda row: row.key)
    return sorted(result, key=lambda row: (row.cost or Decimal(0), row.total_tokens), reverse=True)


def names(session: Session, group: str, keys: set[Any]) -> dict[Any, str]:
    ids = {k for k in keys if isinstance(k, int)}
    labels: dict[Any, str] = {None: "(gelöscht)"}
    if group == "team":
        labels |= {t.id: t.name for t in session.scalars(select(Team).where(Team.id.in_(ids)))}
    elif group == "workspace":
        labels |= {w.id: w.name for w in session.scalars(select(Workspace).where(Workspace.id.in_(ids)))}
    elif group == "user":
        labels |= {u.id: u.display_name for u in session.scalars(select(User).where(User.id.in_(ids)))}
    elif group == "day":
        labels |= {k: k.strftime("%d.%m.%Y") for k in keys if k is not None}
    else:
        labels |= {k: k for k in keys if k is not None}
    return labels


def total(rows: list[Row]) -> Row:
    costs = [r.cost for r in rows if r.cost is not None]
    return Row(None, "Summe", sum(r.calls for r in rows), sum(r.runs for r in rows),
               sum(r.input_tokens for r in rows), sum(r.cached_tokens for r in rows),
               sum(r.output_tokens for r in rows), sum(r.reasoning_tokens for r in rows),
               sum(costs, Decimal(0)) if costs else None, sum(r.unpriced for r in rows))


def run_stats(session: Session, start: datetime, end: datetime, *, team_ids: set[int] | None = None,
              workspace_id: int | None = None) -> dict[str, int]:
    """Runs im Zeitraum je Status."""
    query = (select(TaskRun.status, func.count()).where(TaskRun.created >= start, TaskRun.created < end)
             .group_by(TaskRun.status))
    if team_ids is not None:
        query = query.where(TaskRun.team_id.in_(team_ids))
    if workspace_id is not None:
        query = query.where(TaskRun.workspace_id == workspace_id)
    return dict(session.execute(query).all())


@dataclass
class RunRow:
    run: TaskRun
    workspace: str
    user: str
    tokens: int
    cost: Decimal | None

    @property
    def duration(self) -> float | None:
        if self.run.started and self.run.finished:
            return (self.run.finished - self.run.started).total_seconds()
        return None


def runs(session: Session, start: datetime, end: datetime, *, status: str = "", workspace_id: int | None = None,
         team_ids: set[int] | None = None, limit: int = 200) -> list[RunRow]:
    """Letzte Runs im Zeitraum mit Verbrauch; neueste zuerst."""
    query = select(TaskRun).where(TaskRun.created >= start, TaskRun.created < end)
    if status:
        query = query.where(TaskRun.status == status)
    if workspace_id is not None:
        query = query.where(TaskRun.workspace_id == workspace_id)
    if team_ids is not None:
        query = query.where(TaskRun.team_id.in_(team_ids))
    found = list(session.scalars(query.order_by(TaskRun.created.desc()).limit(limit)))
    ids = [r.id for r in found]
    spent = {row[0]: (int(row[1] or 0), row[2]) for row in session.execute(
        select(UsageRecord.run_id, func.sum(UsageRecord.input_tokens + UsageRecord.output_tokens),
               func.sum(UsageRecord.cost)).where(UsageRecord.run_id.in_(ids)).group_by(UsageRecord.run_id))}
    workspaces = names(session, "workspace", {r.workspace_id for r in found})
    users = names(session, "user", {r.created_by for r in found})
    return [RunRow(r, workspaces.get(r.workspace_id, ""), users.get(r.created_by, ""), *spent.get(r.id, (0, None)))
            for r in found]