"""Message-level data fetching.

Messages are the core unit — every assistant turn, user turn, tool call, and
tool result is a message row.
"""

from datetime import datetime

import pandas as pd
import psycopg

from .connection import query_df


# ---------------------------------------------------------------------------
# By session
# ---------------------------------------------------------------------------

def by_session(
    conn: psycopg.Connection,
    session_id: str,
    *,
    msg_type: str | None = None,
    limit: int = 2000,
) -> pd.DataFrame:
    """All messages for a session, optionally filtered by msg_type."""
    clauses = ["m.session_id = %s"]
    params: list = [session_id]
    if msg_type is not None:
        clauses.append("m.msg_type = %s")
        params.append(msg_type)
    sql = f"""
        SELECT m.*, s.user_email, s.user_name, s.org, s.repo_name
        FROM   messages m
        JOIN   sessions s ON s.id = m.session_id
        WHERE  {' AND '.join(clauses)}
        ORDER  BY m.timestamp ASC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def by_sessions(
    conn: psycopg.Connection,
    session_ids: list[str],
    *,
    msg_type: str | None = None,
    limit: int = 5000,
) -> pd.DataFrame:
    """Messages across multiple sessions."""
    clauses = ["m.session_id = ANY(%s)"]
    params: list = [session_ids]
    if msg_type is not None:
        clauses.append("m.msg_type = %s")
        params.append(msg_type)
    sql = f"""
        SELECT m.*, s.user_email, s.user_name, s.org, s.repo_name
        FROM   messages m
        JOIN   sessions s ON s.id = m.session_id
        WHERE  {' AND '.join(clauses)}
        ORDER  BY m.timestamp ASC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


# ---------------------------------------------------------------------------
# By user / org / time
# ---------------------------------------------------------------------------

def by_user(
    conn: psycopg.Connection,
    email: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    msg_type: str | None = None,
    limit: int = 5000,
) -> pd.DataFrame:
    """Messages for a user across all their sessions."""
    clauses = ["s.user_email = %s"]
    params: list = [email]
    _add_msg_filters(clauses, params, since, until, msg_type)
    return _joined_query(conn, clauses, params, limit)


def by_users(
    conn: psycopg.Connection,
    emails: list[str],
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    msg_type: str | None = None,
    limit: int = 10000,
) -> pd.DataFrame:
    """Messages for multiple users."""
    clauses = ["s.user_email = ANY(%s)"]
    params: list = [emails]
    _add_msg_filters(clauses, params, since, until, msg_type)
    return _joined_query(conn, clauses, params, limit)


def by_org(
    conn: psycopg.Connection,
    org: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    msg_type: str | None = None,
    limit: int = 10000,
) -> pd.DataFrame:
    """Messages for an entire org."""
    clauses = ["s.org = %s"]
    params: list = [org]
    _add_msg_filters(clauses, params, since, until, msg_type)
    return _joined_query(conn, clauses, params, limit)


def by_repo(
    conn: psycopg.Connection,
    repo_name: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    msg_type: str | None = None,
    limit: int = 5000,
) -> pd.DataFrame:
    """Messages for a specific repo."""
    clauses = ["s.repo_name = %s"]
    params: list = [repo_name]
    _add_msg_filters(clauses, params, since, until, msg_type)
    return _joined_query(conn, clauses, params, limit)


def search(
    conn: psycopg.Connection,
    *,
    email: str | None = None,
    org: str | None = None,
    repo_name: str | None = None,
    session_id: str | None = None,
    msg_type: str | None = None,
    source: str | None = None,
    model: str | None = None,
    content_like: str | None = None,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 5000,
) -> pd.DataFrame:
    """Flexible multi-filter message search."""
    clauses: list[str] = []
    params: list = []
    _add_eq(clauses, params, "s.user_email", email)
    _add_eq(clauses, params, "s.org", org)
    _add_eq(clauses, params, "s.repo_name", repo_name)
    _add_eq(clauses, params, "m.session_id", session_id)
    _add_eq(clauses, params, "m.msg_type", msg_type)
    _add_eq(clauses, params, "m.source", source)
    _add_eq(clauses, params, "m.model", model)
    if content_like is not None:
        clauses.append("m.content ILIKE %s")
        params.append(f"%{content_like}%")
    _add_msg_filters(clauses, params, since, until, msg_type=None)
    return _joined_query(conn, clauses, params, limit)


# ---------------------------------------------------------------------------
# Distinct values (useful for filters in notebooks)
# ---------------------------------------------------------------------------

def distinct_msg_types(conn: psycopg.Connection) -> list[str]:
    """All distinct msg_type values."""
    df = query_df(conn, "SELECT DISTINCT msg_type FROM messages ORDER BY msg_type")
    return df["msg_type"].tolist()


def distinct_sources(conn: psycopg.Connection) -> list[str]:
    """All distinct source values."""
    df = query_df(conn, "SELECT DISTINCT source FROM messages ORDER BY source")
    return df["source"].tolist()


def distinct_models(conn: psycopg.Connection) -> list[str]:
    """All distinct model values."""
    df = query_df(conn, "SELECT DISTINCT model FROM messages WHERE model IS NOT NULL ORDER BY model")
    return df["model"].tolist()


# ---------------------------------------------------------------------------
# Counts
# ---------------------------------------------------------------------------

def count_by_type(
    conn: psycopg.Connection,
    session_id: str | None = None,
) -> pd.DataFrame:
    """Message count grouped by msg_type."""
    if session_id:
        return query_df(conn, """
            SELECT msg_type, COUNT(*) AS cnt
            FROM   messages WHERE session_id = %s
            GROUP  BY msg_type ORDER BY cnt DESC
        """, (session_id,))
    return query_df(conn, """
        SELECT msg_type, COUNT(*) AS cnt
        FROM   messages GROUP BY msg_type ORDER BY cnt DESC
    """)


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

def _joined_query(
    conn: psycopg.Connection,
    clauses: list[str],
    params: list,
    limit: int,
) -> pd.DataFrame:
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    sql = f"""
        SELECT m.*, s.user_email, s.user_name, s.org, s.repo_name
        FROM   messages m
        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_msg_filters(
    clauses: list, params: list,
    since: datetime | None, until: datetime | None,
    msg_type: str | 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)
    if msg_type is not None:
        clauses.append("m.msg_type = %s")
        params.append(msg_type)


def _add_eq(clauses: list, params: list, col: str, val) -> None:
    if val is not None:
        clauses.append(f"{col} = %s")
        params.append(val)
