"""Base analyzer and data loading."""

import pandas as pd
import psycopg

from .config import get_connection, query_df


def load_session_data(conn: psycopg.Connection, email: str, since: str, until: str) -> dict:
    """Load all data for a user in a date range. Returns dict of DataFrames."""
    sessions = query_df(conn, """
        SELECT * FROM sessions
        WHERE user_email = %s AND first_seen >= %s AND first_seen < %s
        ORDER BY first_seen ASC
    """, (email, since, until))

    if sessions.empty:
        return {
            "sessions": sessions,
            "messages": pd.DataFrame(),
            "tool_calls": pd.DataFrame(),
            "tool_results": pd.DataFrame(),
            "token_usage": pd.DataFrame(),
        }

    session_ids = sessions["id"].tolist()

    messages = query_df(conn, """
        SELECT m.*
        FROM messages m
        WHERE m.session_id = ANY(%s)
        ORDER BY m.session_id, m.timestamp ASC
    """, (session_ids,))

    message_ids = messages["id"].tolist() if not messages.empty else []

    if message_ids:
        tool_calls = query_df(conn, """
            SELECT tc.*, m.session_id, m.timestamp
            FROM tool_calls tc
            JOIN messages m ON m.id = tc.message_id
            WHERE tc.message_id = ANY(%s)
            ORDER BY m.timestamp ASC
        """, (message_ids,))

        tool_results = query_df(conn, """
            SELECT tr.*, m.session_id, m.timestamp
            FROM tool_results tr
            JOIN messages m ON m.id = tr.message_id
            WHERE tr.message_id = ANY(%s)
            ORDER BY m.timestamp ASC
        """, (message_ids,))

        token_usage = query_df(conn, """
            SELECT t.*, m.session_id, m.msg_type, m.timestamp, m.model
            FROM token_usage t
            JOIN messages m ON m.id = t.message_id
            WHERE t.message_id = ANY(%s)
            ORDER BY m.timestamp ASC
        """, (message_ids,))
    else:
        tool_calls = pd.DataFrame()
        tool_results = pd.DataFrame()
        token_usage = pd.DataFrame()

    return {
        "sessions": sessions,
        "messages": messages,
        "tool_calls": tool_calls,
        "tool_results": tool_results,
        "token_usage": token_usage,
    }


class BaseAnalyzer:
    """Base class for all analyzers.

    Subclasses implement analyze() which receives all DataFrames.

    DB schema reminder:
    - messages.msg_type: 'user', 'assistant', 'tool_call', 'tool_result'
    - messages.content: text content
    - messages.source: 'claude_code', 'codex_cli', 'gemini_cli', 'cursor'
    - tool_calls: message_id, tool_id, tool_name, tool_input (JSONB)
    - tool_results: message_id, call_id, output (TEXT, max 2000 chars), status
    """

    name: str = "base"

    def analyze(self, sessions_df, messages_df, tool_calls_df, tool_results_df, token_usage_df) -> dict:
        raise NotImplementedError

    def get_session_messages(self, messages_df: pd.DataFrame, session_id: str) -> pd.DataFrame:
        """Get messages for a single session, ordered by timestamp."""
        return messages_df[messages_df["session_id"] == session_id].sort_values("timestamp")

    def get_session_tool_calls(self, tool_calls_df: pd.DataFrame, session_id: str) -> pd.DataFrame:
        """Get tool calls for a single session."""
        if tool_calls_df.empty or "session_id" not in tool_calls_df.columns:
            return pd.DataFrame()
        return tool_calls_df[tool_calls_df["session_id"] == session_id]

    def get_session_tool_results(self, tool_results_df: pd.DataFrame, session_id: str) -> pd.DataFrame:
        """Get tool results for a single session."""
        if tool_results_df.empty or "session_id" not in tool_results_df.columns:
            return pd.DataFrame()
        return tool_results_df[tool_results_df["session_id"] == session_id]

    def iter_sessions(self, sessions_df: pd.DataFrame):
        """Iterate over sessions as dicts."""
        for _, row in sessions_df.iterrows():
            yield row.to_dict()

    def build_message_stream(self, messages_df: pd.DataFrame, tool_calls_df: pd.DataFrame,
                              tool_results_df: pd.DataFrame, session_id: str) -> list[dict]:
        """Build a unified message stream for a session, merging tool info into messages.

        Returns list of dicts with keys: id, msg_type, timestamp, content, source,
        tool_name (for tool_call), tool_input (for tool_call),
        result_output (for tool_result), result_status (for tool_result).
        """
        msgs = self.get_session_messages(messages_df, session_id)
        if msgs.empty:
            return []

        # Build lookup: message_id -> tool_call info
        tc_map = {}
        if not tool_calls_df.empty and "session_id" in tool_calls_df.columns:
            stc = tool_calls_df[tool_calls_df["session_id"] == session_id]
            for _, tc in stc.iterrows():
                tn = tc.get("tool_name")
                ti = tc.get("tool_input")
                tc_map[tc["message_id"]] = {
                    "tool_name": tn if isinstance(tn, str) else "",
                    "tool_input": ti if isinstance(ti, (str, dict)) else "",
                }

        # Build lookup: message_id -> tool_result info
        tr_map = {}
        if not tool_results_df.empty and "session_id" in tool_results_df.columns:
            str_ = tool_results_df[tool_results_df["session_id"] == session_id]
            for _, tr in str_.iterrows():
                ro = tr.get("output")
                rs = tr.get("status")
                tr_map[tr["message_id"]] = {
                    "result_output": ro if isinstance(ro, str) else "",
                    "result_status": rs if isinstance(rs, str) else "",
                }

        stream = []
        for _, m in msgs.iterrows():
            entry = {
                "id": m["id"],
                "msg_type": m["msg_type"],
                "timestamp": str(m["timestamp"]) if m.get("timestamp") else None,
                "content": m["content"] if isinstance(m.get("content"), str) else "",
                "source": m["source"] if isinstance(m.get("source"), str) else "",
            }
            mid = m["id"]
            if mid in tc_map:
                entry.update(tc_map[mid])
            if mid in tr_map:
                entry.update(tr_map[mid])
            stream.append(entry)

        return stream
