"""
SQL Trace Persistence — Save/load trace sessions to/from PostgreSQL.
"""

import json
import logging
from typing import Any

from agent.sql_agent.session_manager import public_session
from common.pool_wrapper import get_pooled_connection_compat

logger = logging.getLogger(__name__)


def ensure_sql_tables() -> None:
    """Create sql_trace_sessions table if it doesn't exist."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute("""
            CREATE TABLE IF NOT EXISTS dashboard_canifa.sql_trace_sessions (
                id SERIAL PRIMARY KEY,
                conversation_id UUID NOT NULL,
                session_id INT NOT NULL,
                question TEXT NOT NULL,
                model VARCHAR(100),
                status VARCHAR(20) DEFAULT 'running',
                session_data JSONB,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                completed_at TIMESTAMPTZ
            );
            CREATE INDEX IF NOT EXISTS idx_sql_trace_conv_id
                ON dashboard_canifa.sql_trace_sessions(conversation_id);
        """)
        cur.close()
    except Exception as e:
        logger.error("Error creating sql_trace_sessions table: %s", e)
    finally:
        if conn:
            conn.close()


def persist_session_to_db(session: dict[str, Any]) -> None:
    """Save completed session to PostgreSQL."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        public = public_session(session)
        cur.execute(
            """
            INSERT INTO dashboard_canifa.sql_trace_sessions
                (conversation_id, session_id, question, model, status, session_data, created_at, completed_at)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT DO NOTHING
        """,
            (
                session["conversation_id"],
                session["id"],
                session["question"],
                session.get("model", ""),
                session["status"],
                json.dumps(public, ensure_ascii=False, default=str),
                session["created_at"],
                session.get("completed_at"),
            ),
        )
        cur.close()
        logger.info(
            "Persisted SQL trace session %s (conv=%s)",
            session["id"],
            session["conversation_id"],
        )
    except Exception as e:
        logger.error("Error persisting SQL trace session: %s", e)
    finally:
        if conn:
            conn.close()


def load_conversations(limit: int = 50) -> list[dict[str, Any]]:
    """Get list of SQL trace conversations from DB, newest first."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            """
            SELECT conversation_id,
                   MIN(question) AS title,
                   MAX(status) AS status,
                   MIN(created_at) AS created_at,
                   MAX(completed_at) AS updated_at,
                   COUNT(*) AS session_count
            FROM dashboard_canifa.sql_trace_sessions
            GROUP BY conversation_id
            ORDER BY MAX(created_at) DESC
            LIMIT %s
        """,
            (limit,),
        )
        conversations = []
        for r in cur.fetchall():
            title = r[1] or "Untitled"
            conversations.append(
                {
                    "conversation_id": str(r[0]),
                    "title": title[:80] + "..." if len(title) > 80 else title,
                    "status": r[2],
                    "created_at": r[3].isoformat() if r[3] else None,
                    "updated_at": r[4].isoformat() if r[4] else None,
                    "session_count": r[5],
                }
            )
        cur.close()
        return conversations
    except Exception as e:
        logger.error("Error listing SQL conversations: %s", e)
        return []
    finally:
        if conn:
            conn.close()


def load_conversation_detail(conv_id: str) -> list[dict[str, Any]]:
    """Get all sessions for a specific conversation from DB."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            """
            SELECT session_id, question, model, status, session_data, created_at, completed_at
            FROM dashboard_canifa.sql_trace_sessions
            WHERE conversation_id = %s
            ORDER BY created_at ASC
        """,
            (conv_id,),
        )
        sessions = []
        for r in cur.fetchall():
            session_data = (
                r[4]
                if isinstance(r[4], dict)
                else json.loads(r[4])
                if r[4]
                else {}
            )
            sessions.append(
                {
                    "session_id": r[0],
                    "question": r[1],
                    "model": r[2],
                    "status": r[3],
                    "session_data": session_data,
                    "created_at": r[5].isoformat() if r[5] else None,
                    "completed_at": r[6].isoformat() if r[6] else None,
                }
            )
        cur.close()
        return sessions
    except Exception as e:
        logger.error("Error fetching SQL conversation %s: %s", conv_id, e)
        return []
    finally:
        if conn:
            conn.close()


# Auto-init tables on import
try:
    ensure_sql_tables()
except Exception:
    pass
