"""Tool call and tool result data fetching."""

from datetime import datetime

import pandas as pd
import psycopg

from .connection import query_df


# ---------------------------------------------------------------------------
# Tool calls
# ---------------------------------------------------------------------------

def calls_by_session(conn: psycopg.Connection, session_id: str) -> pd.DataFrame:
    """Tool calls for a session."""
    return query_df(conn, """
        SELECT tc.*, m.timestamp, m.model
        FROM   tool_calls tc
        JOIN   messages m ON m.id = tc.message_id
        WHERE  m.session_id = %s
        ORDER  BY m.timestamp ASC
    """, (session_id,))


def calls_by_user(
    conn: psycopg.Connection,
    email: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    tool_name: str | None = None,
    limit: int = 5000,
) -> pd.DataFrame:
    """Tool calls for a user."""
    clauses = ["s.user_email = %s"]
    params: list = [email]
    _add_filters(clauses, params, since, until, tool_name)
    return _calls_joined(conn, clauses, params, limit)


def calls_by_org(
    conn: psycopg.Connection,
    org: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    tool_name: str | None = None,
    limit: int = 10000,
) -> pd.DataFrame:
    """Tool calls for an org."""
    clauses = ["s.org = %s"]
    params: list = [org]
    _add_filters(clauses, params, since, until, tool_name)
    return _calls_joined(conn, clauses, params, limit)


# ---------------------------------------------------------------------------
# Tool results
# ---------------------------------------------------------------------------

def results_by_session(conn: psycopg.Connection, session_id: str) -> pd.DataFrame:
    """Tool results for a session."""
    return query_df(conn, """
        SELECT tr.*, m.timestamp
        FROM   tool_results tr
        JOIN   messages m ON m.id = tr.message_id
        WHERE  m.session_id = %s
        ORDER  BY m.timestamp ASC
    """, (session_id,))


def results_by_user(
    conn: psycopg.Connection,
    email: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    status: str | None = None,
    limit: int = 5000,
) -> pd.DataFrame:
    """Tool results for a user, optionally filtered by status."""
    clauses = ["s.user_email = %s"]
    params: list = [email]
    _add_time(clauses, params, since, until)
    if status is not None:
        clauses.append("tr.status = %s")
        params.append(status)
    return _results_joined(conn, clauses, params, limit)


# ---------------------------------------------------------------------------
# Aggregations
# ---------------------------------------------------------------------------

def call_frequency(
    conn: psycopg.Connection,
    *,
    email: str | None = None,
    org: str | None = None,
    session_id: str | None = None,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Tool usage frequency: how many times each tool was called."""
    clauses: list[str] = []
    params: list = []
    needs_session = email is not None or org is not None
    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)
    if session_id is not None:
        clauses.append("m.session_id = %s")
        params.append(session_id)
    _add_time(clauses, params, since, until)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    join_sessions = "JOIN sessions s ON s.id = m.session_id" if needs_session else ""
    return query_df(conn, f"""
        SELECT tc.tool_name,
               COUNT(*)               AS call_count,
               COUNT(DISTINCT m.session_id) AS session_count
        FROM   tool_calls tc
        JOIN   messages m ON m.id = tc.message_id
        {join_sessions}
        {where}
        GROUP  BY tc.tool_name
        ORDER  BY call_count DESC
    """, tuple(params))


def failure_rate(
    conn: psycopg.Connection,
    *,
    email: str | None = None,
    org: str | None = None,
    since: datetime | None = None,
    until: datetime | None = None,
) -> pd.DataFrame:
    """Per-tool success/failure counts. Joins tool_calls with tool_results."""
    clauses: list[str] = []
    params: list = []
    needs_session = email is not None or org is not None
    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 ""
    join_sessions = "JOIN sessions s ON s.id = m.session_id" if needs_session else ""
    return query_df(conn, f"""
        SELECT tc.tool_name,
               COUNT(*)                                      AS total,
               COUNT(*) FILTER (WHERE tr.status = 'success') AS successes,
               COUNT(*) FILTER (WHERE tr.status = 'failure') AS failures
        FROM   tool_calls tc
        JOIN   messages m ON m.id = tc.message_id
        LEFT JOIN tool_results tr ON tr.call_id = tc.tool_id
                                  AND tr.message_id IN (
                                      SELECT id FROM messages
                                      WHERE session_id = m.session_id
                                  )
        {join_sessions}
        {where}
        GROUP  BY tc.tool_name
        ORDER  BY failures DESC
    """, tuple(params))


def distinct_tool_names(conn: psycopg.Connection) -> list[str]:
    """All distinct tool names."""
    df = query_df(conn, "SELECT DISTINCT tool_name FROM tool_calls ORDER BY tool_name")
    return df["tool_name"].tolist()


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

def _calls_joined(
    conn: psycopg.Connection,
    clauses: list[str],
    params: list,
    limit: int,
) -> pd.DataFrame:
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    sql = f"""
        SELECT tc.*, m.session_id, m.timestamp, m.model,
               s.user_email, s.org, s.repo_name
        FROM   tool_calls tc
        JOIN   messages m ON m.id = tc.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 _results_joined(
    conn: psycopg.Connection,
    clauses: list[str],
    params: list,
    limit: int,
) -> pd.DataFrame:
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    sql = f"""
        SELECT tr.*, m.session_id, m.timestamp,
               s.user_email, s.org, s.repo_name
        FROM   tool_results tr
        JOIN   messages m ON m.id = tr.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_filters(
    clauses: list, params: list,
    since: datetime | None, until: datetime | None,
    tool_name: str | None,
) -> None:
    _add_time(clauses, params, since, until)
    if tool_name is not None:
        clauses.append("tc.tool_name = %s")
        params.append(tool_name)


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)
