"""
SQL Trace Session Manager — In-memory session state management.

Manages trace session lifecycle: create, update cycles, record events, SSE formatting.
Mirrors the pattern used in report_agent/report_queue.py.
"""

import asyncio
import json
import time
import uuid
from datetime import datetime, timezone
from itertools import count
from typing import Any

_TRACE_SESSION_SEQ = count(1)
_TRACE_SESSIONS: dict[int, dict[str, Any]] = {}
_MAX_STORED_SESSIONS = 80

# Human-in-the-Loop: approval events keyed by session_id
_APPROVAL_EVENTS: dict[int, asyncio.Event] = {}
_APPROVAL_ACTIONS: dict[int, str] = {}  # "accept" or "reject"


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def prune_sessions() -> None:
    """Remove oldest sessions when exceeding max capacity."""
    if len(_TRACE_SESSIONS) <= _MAX_STORED_SESSIONS:
        return

    ordered_ids = sorted(
        _TRACE_SESSIONS.keys(),
        key=lambda sid: _TRACE_SESSIONS[sid].get("_created_ts", 0.0),
    )
    while len(ordered_ids) > _MAX_STORED_SESSIONS:
        sid = ordered_ids.pop(0)
        _TRACE_SESSIONS.pop(sid, None)


def new_cycle(cycle_no: int) -> dict[str, Any]:
    """Create a fresh cycle object."""
    return {
        "cycle": cycle_no,
        "thinking": "",
        "tool_calls": [],
        "tool_results": [],
        "missing": [],
        "data_sufficient": False,
        "reflect_thinking": "",
        "status": "pending",
    }


def create_session(
    question: str,
    model: str,
    max_cycles: int,
    conversation_id: str | None = None,
) -> dict[str, Any]:
    """Create a new trace session and register it."""
    session_id = next(_TRACE_SESSION_SEQ)
    now = utc_now_iso()
    conv_id = conversation_id or str(uuid.uuid4())
    session = {
        "id": session_id,
        "conversation_id": conv_id,
        "question": question,
        "model": model,
        "max_cycles": max_cycles,
        "status": "running",
        "created_at": now,
        "updated_at": now,
        "completed_at": None,
        "events": [],
        "cycles": [],
        "summary": {},
        "total_queries": 0,
        "successful_queries": 0,
        "failed_queries": 0,
        "_created_ts": time.time(),
    }
    _TRACE_SESSIONS[session_id] = session
    prune_sessions()
    return session


# ─── Human-in-the-Loop Approval ─────────────────────────────────────

def create_approval_event(session_id: int) -> asyncio.Event:
    """Create new approval event for a session (resets if exists)."""
    evt = asyncio.Event()
    _APPROVAL_EVENTS[session_id] = evt
    _APPROVAL_ACTIONS[session_id] = ""
    return evt


def approve_session(session_id: int) -> bool:
    """Accept the pending SQL. Returns False if no event found."""
    evt = _APPROVAL_EVENTS.get(session_id)
    if not evt:
        return False
    _APPROVAL_ACTIONS[session_id] = "accept"
    evt.set()
    return True


def reject_session(session_id: int) -> bool:
    """Reject the pending SQL. Returns False if no event found."""
    evt = _APPROVAL_EVENTS.get(session_id)
    if not evt:
        return False
    _APPROVAL_ACTIONS[session_id] = "reject"
    evt.set()
    return True


def get_approval_action(session_id: int) -> str:
    """Get approval action (accept/reject) for a session."""
    return _APPROVAL_ACTIONS.get(session_id, "")


def cleanup_approval(session_id: int) -> None:
    """Remove approval state for a session."""
    _APPROVAL_EVENTS.pop(session_id, None)
    _APPROVAL_ACTIONS.pop(session_id, None)


def public_session(session: dict[str, Any]) -> dict[str, Any]:
    """Return session dict without internal keys."""
    return {k: v for k, v in session.items() if not k.startswith("_")}


def ensure_cycle(session: dict[str, Any], cycle_no: int) -> dict[str, Any]:
    """Ensure cycle_no exists in session, creating missing cycles if needed."""
    while len(session["cycles"]) < cycle_no:
        session["cycles"].append(new_cycle(len(session["cycles"]) + 1))
    return session["cycles"][cycle_no - 1]


def record_event(session: dict[str, Any], payload: dict[str, Any]) -> None:
    """Record an event in the session event log."""
    session["updated_at"] = utc_now_iso()
    session["events"].append({"ts": session["updated_at"], **payload})
    if len(session["events"]) > 800:
        session["events"] = session["events"][-800:]


def sse(session: dict[str, Any], payload: dict[str, Any]) -> str:
    """Format a Server-Sent Event string and record it."""
    data = {"session_id": session["id"], **payload}
    record_event(session, data)
    return f"data: {json.dumps(data, ensure_ascii=False, default=str)}\n\n"


def get_session(session_id: int) -> dict[str, Any] | None:
    """Get session by ID."""
    return _TRACE_SESSIONS.get(session_id)


def delete_session(session_id: int) -> bool:
    """Delete session by ID. Returns True if found."""
    if session_id in _TRACE_SESSIONS:
        del _TRACE_SESSIONS[session_id]
        return True
    return False


def list_sessions() -> list[dict[str, Any]]:
    """List all sessions (summary info only)."""
    sessions = []
    for session in _TRACE_SESSIONS.values():
        sessions.append(
            {
                "id": session["id"],
                "question": session["question"],
                "model": session["model"],
                "status": session["status"],
                "created_at": session["created_at"],
                "updated_at": session["updated_at"],
                "completed_at": session["completed_at"],
                "total_queries": session["total_queries"],
                "successful_queries": session["successful_queries"],
                "failed_queries": session["failed_queries"],
                "cycles_completed": session.get("summary", {}).get(
                    "cycles_completed", len(session["cycles"])
                ),
                "data_sufficient": session.get("summary", {}).get("data_sufficient"),
                "final_task": session.get("summary", {}).get("final_task", ""),
            }
        )
    sessions.sort(key=lambda x: x.get("created_at") or "", reverse=True)
    return sessions
