"""Token usage data fetching.

Joins token_usage → messages → sessions so every query can filter by user,
org, repo, or time range.
"""

from datetime import datetime

import pandas as pd
import psycopg

from .connection import query_df


# ---------------------------------------------------------------------------
# Raw token rows
# ---------------------------------------------------------------------------

def by_session(conn: psycopg.Connection, session_id: str) -> pd.DataFrame:
    """Token usage for every message in a session."""
    return query_df(conn, """
        SELECT t.*, m.msg_type, m.timestamp, m.model
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        WHERE  m.session_id = %s
        ORDER  BY m.timestamp ASC
    """, (session_id,))


def by_user(
    conn: psycopg.Connection,
    email: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Token rows for a user."""
    clauses = ["s.user_email = %s"]
    params: list = [email]
    _add_time(clauses, params, since, until)
    return _joined(conn, clauses, params)


def by_users(
    conn: psycopg.Connection,
    emails: list[str],
    *,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Token rows for multiple users."""
    clauses = ["s.user_email = ANY(%s)"]
    params: list = [emails]
    _add_time(clauses, params, since, until)
    return _joined(conn, clauses, params)


def by_org(
    conn: psycopg.Connection,
    org: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Token rows for an org."""
    clauses = ["s.org = %s"]
    params: list = [org]
    _add_time(clauses, params, since, until)
    return _joined(conn, clauses, params)


# ---------------------------------------------------------------------------
# Aggregated summaries
# ---------------------------------------------------------------------------

def summary_by_session(
    conn: psycopg.Connection,
    session_ids: list[str] | None = None,
) -> pd.DataFrame:
    """Per-session totals: input, output, cached, thinking tokens."""
    where = "WHERE m.session_id = ANY(%s)" if session_ids else ""
    params = (session_ids,) if session_ids else ()
    return query_df(conn, f"""
        SELECT m.session_id,
               SUM(t.input_tokens)    AS total_input,
               SUM(t.output_tokens)   AS total_output,
               SUM(t.cached_tokens)   AS total_cached,
               SUM(t.thinking_tokens) AS total_thinking,
               SUM(t.input_tokens + t.output_tokens) AS total_billed,
               COUNT(*)               AS message_count
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        {where}
        GROUP  BY m.session_id
        ORDER  BY total_billed DESC
    """, params)


def summary_by_user(
    conn: psycopg.Connection,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Per-user token totals."""
    clauses: list[str] = []
    params: list = []
    _add_time(clauses, params, since, until)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    return query_df(conn, f"""
        SELECT s.user_email, s.org,
               SUM(t.input_tokens)    AS total_input,
               SUM(t.output_tokens)   AS total_output,
               SUM(t.cached_tokens)   AS total_cached,
               SUM(t.thinking_tokens) AS total_thinking,
               SUM(t.input_tokens + t.output_tokens) AS total_billed,
               COUNT(DISTINCT m.session_id) AS session_count,
               COUNT(*)               AS message_count
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        JOIN   sessions s ON s.id = m.session_id
        {where}
        GROUP  BY s.user_email, s.org
        ORDER  BY total_billed DESC
    """, tuple(params))


def summary_by_org(
    conn: psycopg.Connection,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Per-org token totals."""
    clauses: list[str] = []
    params: list = []
    _add_time(clauses, params, since, until)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    return query_df(conn, f"""
        SELECT s.org,
               SUM(t.input_tokens)    AS total_input,
               SUM(t.output_tokens)   AS total_output,
               SUM(t.cached_tokens)   AS total_cached,
               SUM(t.thinking_tokens) AS total_thinking,
               SUM(t.input_tokens + t.output_tokens) AS total_billed,
               COUNT(DISTINCT s.user_email) AS user_count,
               COUNT(DISTINCT m.session_id) AS session_count,
               COUNT(*)               AS message_count
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        JOIN   sessions s ON s.id = m.session_id
        {where}
        GROUP  BY s.org
        ORDER  BY total_billed DESC
    """, tuple(params))


def summary_by_model(
    conn: psycopg.Connection,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Per-model token totals."""
    clauses: list[str] = []
    params: list = []
    _add_time(clauses, params, since, until)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    return query_df(conn, f"""
        SELECT m.model,
               SUM(t.input_tokens)    AS total_input,
               SUM(t.output_tokens)   AS total_output,
               SUM(t.cached_tokens)   AS total_cached,
               SUM(t.thinking_tokens) AS total_thinking,
               SUM(t.input_tokens + t.output_tokens) AS total_billed,
               COUNT(*)               AS message_count
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        {where}
        GROUP  BY m.model
        ORDER  BY total_billed DESC
    """, tuple(params))


def summary_by_day(
    conn: psycopg.Connection,
    *,
    email: str | None = None,
    org: str | None = None,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Daily token totals, optionally scoped to user or org."""
    clauses: list[str] = []
    params: list = []
    if email is not None:
        clauses.append("s.user_email = %s")
        params.append(email)
    if org is not None:
        clauses.append("s.org = %s")
        params.append(org)
    _add_time(clauses, params, since, until)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    return query_df(conn, f"""
        SELECT DATE(m.timestamp) AS day,
               SUM(t.input_tokens)    AS total_input,
               SUM(t.output_tokens)   AS total_output,
               SUM(t.cached_tokens)   AS total_cached,
               SUM(t.thinking_tokens) AS total_thinking,
               SUM(t.input_tokens + t.output_tokens) AS total_billed,
               COUNT(DISTINCT m.session_id) AS session_count,
               COUNT(*)               AS message_count
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        JOIN   sessions s ON s.id = m.session_id
        {where}
        GROUP  BY day
        ORDER  BY day DESC
    """, tuple(params))


# ---------------------------------------------------------------------------
# Internal
# ---------------------------------------------------------------------------

def _joined(
    conn: psycopg.Connection,
    clauses: list[str],
    params: list,
    limit: int = 10000,
) -> pd.DataFrame:
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    sql = f"""
        SELECT t.*, m.session_id, m.msg_type, m.timestamp, m.model,
               s.user_email, s.org, s.repo_name
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        JOIN   sessions s ON s.id = m.session_id
        {where}
        ORDER  BY m.timestamp ASC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def _add_time(
    clauses: list, params: list,
    since: datetime | None, until: datetime | None,
) -> None:
    if since is not None:
        clauses.append("m.timestamp >= %s")
        params.append(since)
    if until is not None:
        clauses.append("m.timestamp < %s")
        params.append(until)
