"""
Session Tracker — Persistent report lifecycle tracking.

Adapted from hermes-agent-repo/hermes_state.py SessionDB patterns.
Provides high-level lifecycle management for report generation sessions,
wrapping the lower-level ReportSessionDB with convenience methods.

Usage:
    from agent.report_agent.session_tracker import ReportSessionTracker

    tracker = ReportSessionTracker()

    # Start a session
    sid = tracker.start("Doanh thu tháng 5 theo chi nhánh")

    # Track events during generation
    tracker.log_tool_call(sid, "sql_query", {"query": "SELECT..."})
    tracker.log_llm_call(sid, input_tokens=500, output_tokens=200)
    tracker.log_error(sid, error_type="rate_limit", message="429")
    tracker.log_compression(sid)

    # End with final metrics
    tracker.finish(sid, report_id=42, html_length=15000)
"""

from __future__ import annotations

import logging
import time
import uuid
from typing import Any, Dict, List, Optional

from agent.report_agent.hermes_bridge import ReportSessionDB

logger = logging.getLogger(__name__)


class ReportSessionTracker:
    """High-level lifecycle tracking for report generation sessions.

    Each report generation request creates a session that tracks:
      - Question asked
      - Tools invoked and their results
      - LLM calls with token counts
      - Errors encountered and recovery actions
      - Context compressions triggered
      - Final output metrics (HTML size, generation time)
    """

    def __init__(self, db: ReportSessionDB = None):
        self._db = db
        self._active_sessions: Dict[str, Dict[str, Any]] = {}

    @property
    def db(self) -> ReportSessionDB:
        if self._db is None:
            self._db = ReportSessionDB()
        return self._db

    def start(
        self,
        question: str,
        model: str = None,
        session_id: str = None,
    ) -> str:
        """Start tracking a new report generation session.

        Returns the session ID.
        """
        sid = session_id or uuid.uuid4().hex[:12]

        self.db.create_session(sid, question, model=model)
        self._active_sessions[sid] = {
            "started_at": time.time(),
            "tools_used": [],
            "input_tokens": 0,
            "output_tokens": 0,
            "error_count": 0,
            "compression_count": 0,
            "cycles": 0,
        }

        logger.info("Report session started: %s — %s", sid, question[:80])
        return sid

    def log_tool_call(
        self,
        session_id: str,
        tool_name: str,
        tool_data: Dict[str, Any] = None,
    ) -> None:
        """Log a tool invocation within a session."""
        state = self._active_sessions.get(session_id, {})
        if tool_name not in state.get("tools_used", []):
            state.setdefault("tools_used", []).append(tool_name)

        self.db.log_event(session_id, "tool_call", {
            "tool": tool_name,
            **(tool_data or {}),
        })

    def log_llm_call(
        self,
        session_id: str,
        input_tokens: int = 0,
        output_tokens: int = 0,
        model: str = None,
    ) -> None:
        """Log an LLM call with token counts."""
        state = self._active_sessions.get(session_id, {})
        state["input_tokens"] = state.get("input_tokens", 0) + input_tokens
        state["output_tokens"] = state.get("output_tokens", 0) + output_tokens
        state["cycles"] = state.get("cycles", 0) + 1

        self.db.log_event(session_id, "llm_call", {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "model": model,
        })

    def log_error(
        self,
        session_id: str,
        error_type: str,
        message: str = "",
    ) -> None:
        """Log an error that occurred during generation."""
        state = self._active_sessions.get(session_id, {})
        state["error_count"] = state.get("error_count", 0) + 1

        self.db.log_event(session_id, "error", {
            "type": error_type,
            "message": message[:500],
        })
        logger.warning(
            "Report session %s error: %s — %s",
            session_id, error_type, message[:200],
        )

    def log_compression(self, session_id: str) -> None:
        """Log that context compression was triggered."""
        state = self._active_sessions.get(session_id, {})
        state["compression_count"] = state.get("compression_count", 0) + 1

        self.db.log_event(session_id, "compression", {})

    def log_cycle(self, session_id: str, cycle_data: Dict[str, Any] = None) -> None:
        """Log a reflection/execution cycle completion."""
        state = self._active_sessions.get(session_id, {})
        state["cycles"] = state.get("cycles", 0) + 1

        self.db.log_event(session_id, "cycle", cycle_data or {})

    def finish(
        self,
        session_id: str,
        *,
        report_id: int = None,
        html_length: int = 0,
        status: str = "done",
    ) -> Dict[str, Any]:
        """Finish a session and persist final metrics.

        Returns a summary dict of the session.
        """
        state = self._active_sessions.pop(session_id, {})
        started_at = state.get("started_at", time.time())
        generation_time_ms = int((time.time() - started_at) * 1000)

        self.db.end_session(
            session_id,
            status=status,
            cycles_count=state.get("cycles", 0),
            tools_used=state.get("tools_used", []),
            input_tokens=state.get("input_tokens", 0),
            output_tokens=state.get("output_tokens", 0),
            error_count=state.get("error_count", 0),
            compression_count=state.get("compression_count", 0),
            report_id=report_id,
            generation_time_ms=generation_time_ms,
            html_length=html_length,
        )

        # Record aggregate metrics
        self.db.record_metric("generation_time_ms", generation_time_ms)
        self.db.record_metric("html_length", html_length)
        self.db.record_metric(
            "total_tokens",
            state.get("input_tokens", 0) + state.get("output_tokens", 0),
        )

        summary = {
            "session_id": session_id,
            "status": status,
            "generation_time_ms": generation_time_ms,
            "cycles": state.get("cycles", 0),
            "tools_used": state.get("tools_used", []),
            "total_tokens": (
                state.get("input_tokens", 0) + state.get("output_tokens", 0)
            ),
            "errors": state.get("error_count", 0),
            "compressions": state.get("compression_count", 0),
        }

        logger.info(
            "Report session %s finished: status=%s, %dms, %d cycles, %d tokens",
            session_id, status, generation_time_ms,
            summary["cycles"], summary["total_tokens"],
        )
        return summary

    def fail(self, session_id: str, error: str = "") -> Dict[str, Any]:
        """Mark a session as failed."""
        if error:
            self.log_error(session_id, "fatal", error)
        return self.finish(session_id, status="error")

    # ── Query API ──

    def get_recent(self, limit: int = 10) -> List[Dict[str, Any]]:
        """Get recent sessions for dashboard display."""
        return self.db.get_sessions(limit=limit)

    def get_overview(self, days: int = 30) -> Dict[str, Any]:
        """Get aggregate statistics for the last N days."""
        return self.db.get_overview(days=days)

    def get_active_count(self) -> int:
        """Number of sessions currently being tracked in-memory."""
        return len(self._active_sessions)
