"""Session storage — JSON serialization + SQLite index for session records."""

from __future__ import annotations

import json
import logging
import sqlite3
from datetime import datetime
from pathlib import Path

from nestdx.session.models import (
    FileChange,
    GitSnapshot,
    PolicyViolation,
    Session,
    SessionStatus,
    ToolRun,
)
from nestdx.tracking.cost import TokenUsage

logger = logging.getLogger(__name__)

_SCHEMA_SQL = """\
CREATE TABLE IF NOT EXISTS sessions (
    id TEXT PRIMARY KEY,
    status TEXT NOT NULL,
    start_time TEXT NOT NULL,
    end_time TEXT,
    duration_seconds REAL,
    tool_names TEXT,
    file_changes_count INTEGER DEFAULT 0,
    violations_count INTEGER DEFAULT 0,
    total_tokens INTEGER DEFAULT 0,
    total_cost REAL DEFAULT 0.0,
    config_hash TEXT
);
CREATE INDEX IF NOT EXISTS idx_sessions_start_time ON sessions(start_time);
CREATE INDEX IF NOT EXISTS idx_sessions_status ON sessions(status);
"""


def _serialize_datetime(dt: datetime | None) -> str | None:
    return dt.isoformat() if dt else None


def _parse_datetime(s: str | None) -> datetime | None:
    return datetime.fromisoformat(s) if s else None


def session_to_dict(session: Session) -> dict:
    """Convert a Session to a JSON-serializable dict."""
    return {
        "id": session.id,
        "status": session.status.value,
        "start_time": _serialize_datetime(session.start_time),
        "end_time": _serialize_datetime(session.end_time),
        "git_snapshot": (
            {
                "head_sha": session.git_snapshot.head_sha,
                "branch": session.git_snapshot.branch,
                "dirty_files": session.git_snapshot.dirty_files,
            }
            if session.git_snapshot
            else None
        ),
        "tool_runs": [
            {
                "tool": tr.tool,
                "command": tr.command,
                "start_time": _serialize_datetime(tr.start_time),
                "end_time": _serialize_datetime(tr.end_time),
                "exit_code": tr.exit_code,
            }
            for tr in session.tool_runs
        ],
        "file_changes": [
            {
                "path": fc.path,
                "change_type": fc.change_type,
                "additions": fc.additions,
                "deletions": fc.deletions,
            }
            for fc in session.file_changes
        ],
        "violations": [
            {
                "rule": v.rule,
                "severity": v.severity,
                "message": v.message,
                "file": v.file,
                "line": v.line,
                "pattern": v.pattern,
            }
            for v in session.violations
        ],
        "config_hash": session.config_hash,
        "token_usage": [
            {
                "tool": tu.tool,
                "input_tokens": tu.input_tokens,
                "output_tokens": tu.output_tokens,
                "total_tokens": tu.total_tokens,
                "estimated_cost_usd": tu.estimated_cost_usd,
                "model": tu.model,
                "timestamp": _serialize_datetime(tu.timestamp),
            }
            for tu in session.token_usage
        ],
    }


def dict_to_session(data: dict) -> Session:
    """Reconstruct a Session from a deserialized dict."""
    git_snap = None
    if data.get("git_snapshot"):
        gs = data["git_snapshot"]
        git_snap = GitSnapshot(
            head_sha=gs["head_sha"],
            branch=gs["branch"],
            dirty_files=gs.get("dirty_files", []),
        )

    tool_runs = [
        ToolRun(
            tool=tr["tool"],
            command=tr["command"],
            start_time=_parse_datetime(tr["start_time"]),
            end_time=_parse_datetime(tr.get("end_time")),
            exit_code=tr.get("exit_code"),
        )
        for tr in data.get("tool_runs", [])
    ]

    file_changes = [
        FileChange(
            path=fc["path"],
            change_type=fc["change_type"],
            additions=fc.get("additions", 0),
            deletions=fc.get("deletions", 0),
        )
        for fc in data.get("file_changes", [])
    ]

    violations = [
        PolicyViolation(
            rule=v["rule"],
            severity=v["severity"],
            message=v["message"],
            file=v.get("file"),
            line=v.get("line"),
            pattern=v.get("pattern"),
        )
        for v in data.get("violations", [])
    ]

    token_usage_list = []
    for tu in data.get("token_usage", []):
        tu_kwargs = {
            "tool": tu["tool"],
            "input_tokens": tu["input_tokens"],
            "output_tokens": tu["output_tokens"],
            "total_tokens": tu["total_tokens"],
            "estimated_cost_usd": tu["estimated_cost_usd"],
            "model": tu["model"],
        }
        ts = _parse_datetime(tu.get("timestamp"))
        if ts is not None:
            tu_kwargs["timestamp"] = ts
        token_usage_list.append(TokenUsage(**tu_kwargs))

    return Session(
        id=data["id"],
        status=SessionStatus(data["status"]),
        start_time=_parse_datetime(data["start_time"]),
        end_time=_parse_datetime(data.get("end_time")),
        git_snapshot=git_snap,
        tool_runs=tool_runs,
        file_changes=file_changes,
        violations=violations,
        config_hash=data.get("config_hash"),
        token_usage=token_usage_list,
    )


def save_session(session: Session, sessions_dir: Path) -> Path:
    """Write session to JSON file. Returns the file path."""
    sessions_dir.mkdir(parents=True, exist_ok=True)
    path = sessions_dir / f"{session.id}.json"
    path.write_text(json.dumps(session_to_dict(session), indent=2))
    return path


def load_session(path: Path) -> Session | None:
    """Load a session from a JSON file. Returns None on corruption."""
    try:
        data = json.loads(path.read_text())
        return dict_to_session(data)
    except (json.JSONDecodeError, KeyError, ValueError, TypeError) as e:
        logger.warning("Skipping corrupted session file %s: %s", path, e)
        return None


def list_sessions(sessions_dir: Path) -> list[Session]:
    """List all sessions from the sessions directory, newest first."""
    if not sessions_dir.is_dir():
        return []

    sessions = []
    for path in sorted(sessions_dir.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True):
        session = load_session(path)
        if session is not None:
            sessions.append(session)
    return sessions


class SessionStore:
    """Dual storage: JSON files (source of truth) + SQLite (queryable index).

    JSON files remain the canonical session format. SQLite provides fast
    queries, filtering, and aggregation without loading every file.
    """

    def __init__(self, sessions_dir: Path) -> None:
        self.sessions_dir = sessions_dir
        self.sessions_dir.mkdir(parents=True, exist_ok=True)
        self._db_path = sessions_dir / "sessions.db"
        self._init_db()

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(str(self._db_path))
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self) -> None:
        """Create schema and backfill from existing JSON files if needed."""
        needs_backfill = not self._db_path.exists()
        conn = self._connect()
        try:
            conn.executescript(_SCHEMA_SQL)
            if needs_backfill:
                self._backfill(conn)
            conn.commit()
        finally:
            conn.close()

    def _backfill(self, conn: sqlite3.Connection) -> None:
        """Index all existing JSON session files into SQLite."""
        for path in self.sessions_dir.glob("*.json"):
            session = load_session(path)
            if session is not None:
                self._upsert(conn, session)

    def _upsert(self, conn: sqlite3.Connection, session: Session) -> None:
        """Insert or replace a session row in SQLite."""
        tool_names = ",".join(sorted({tr.tool for tr in session.tool_runs}))
        conn.execute(
            """\
            INSERT OR REPLACE INTO sessions
                (id, status, start_time, end_time, duration_seconds,
                 tool_names, file_changes_count, violations_count,
                 total_tokens, total_cost, config_hash)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                session.id,
                session.status.value,
                _serialize_datetime(session.start_time),
                _serialize_datetime(session.end_time),
                session.duration_seconds,
                tool_names,
                len(session.file_changes),
                len(session.violations),
                session.total_tokens,
                session.total_cost,
                session.config_hash,
            ),
        )

    def save(self, session: Session) -> Path:
        """Write session to JSON and upsert into SQLite index.

        Raises:
            OSError: If the disk is full or the write fails.
        """
        try:
            path = save_session(session, self.sessions_dir)
        except OSError as e:
            logger.error("Failed to save session %s: %s", session.id, e)
            raise

        try:
            conn = self._connect()
            try:
                self._upsert(conn, session)
                conn.commit()
            finally:
                conn.close()
        except sqlite3.OperationalError as e:
            # SQLite write failure (e.g. disk full, locked) — JSON is saved,
            # the index can be rebuilt later.
            logger.warning("SQLite index update failed for %s: %s", session.id, e)

        return path

    def load(self, session_id: str) -> Session | None:
        """Load a full session from its JSON file."""
        path = self.sessions_dir / f"{session_id}.json"
        if not path.is_file():
            return None
        return load_session(path)

    def query(
        self,
        *,
        limit: int = 20,
        tool: str | None = None,
        since: datetime | None = None,
        has_violations: bool | None = None,
    ) -> list[Session]:
        """Query sessions via SQLite index, return full Session objects from JSON.

        Args:
            limit: Max sessions to return.
            tool: Filter to sessions that used this tool.
            since: Only sessions started on or after this datetime.
            has_violations: True = only sessions with violations,
                            False = only clean sessions, None = all.
        """
        clauses: list[str] = []
        params: list[object] = []

        if tool is not None:
            clauses.append("(tool_names = ? OR tool_names LIKE ? OR tool_names LIKE ? OR tool_names LIKE ?)")
            params.extend([tool, f"{tool},%", f"%,{tool},%", f"%,{tool}"])

        if since is not None:
            clauses.append("start_time >= ?")
            params.append(_serialize_datetime(since))

        if has_violations is True:
            clauses.append("violations_count > 0")
        elif has_violations is False:
            clauses.append("violations_count = 0")

        where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        sql = f"SELECT id FROM sessions {where} ORDER BY start_time DESC LIMIT ?"
        params.append(limit)

        try:
            conn = self._connect()
            try:
                rows = conn.execute(sql, params).fetchall()
            finally:
                conn.close()
        except sqlite3.OperationalError as e:
            logger.warning("SQLite query failed, falling back to JSON scan: %s", e)
            # Fall back to loading all JSON files
            return list_sessions(self.sessions_dir)[:limit]

        sessions = []
        for row in rows:
            session = self.load(row["id"])
            if session is not None:
                sessions.append(session)
        return sessions

    def aggregate(
        self,
        *,
        group_by: str | None = None,
        since: datetime | None = None,
    ) -> dict:
        """Aggregate session stats from SQLite.

        Args:
            group_by: "tool" or "status". None for overall totals.
            since: Only include sessions on or after this datetime.

        Returns:
            A dict with keys: total_sessions, total_duration, total_tokens,
            total_cost, violations_count, and optionally "groups" if group_by
            is specified.
        """
        where_clauses: list[str] = []
        params: list[object] = []

        if since is not None:
            where_clauses.append("start_time >= ?")
            params.append(_serialize_datetime(since))

        where = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

        conn = self._connect()
        try:
            # Overall totals
            row = conn.execute(
                f"""\
                SELECT
                    COUNT(*) as total_sessions,
                    COALESCE(SUM(duration_seconds), 0) as total_duration,
                    COALESCE(SUM(total_tokens), 0) as total_tokens,
                    COALESCE(SUM(total_cost), 0) as total_cost,
                    COALESCE(SUM(violations_count), 0) as violations_count,
                    COALESCE(SUM(file_changes_count), 0) as file_changes_count
                FROM sessions {where}
                """,
                params,
            ).fetchone()

            result: dict = {
                "total_sessions": row["total_sessions"],
                "total_duration": row["total_duration"],
                "total_tokens": row["total_tokens"],
                "total_cost": row["total_cost"],
                "violations_count": row["violations_count"],
                "file_changes_count": row["file_changes_count"],
            }

            if group_by == "tool":
                result["groups"] = self._aggregate_by_tool(conn, where, params)
            elif group_by == "status":
                result["groups"] = self._aggregate_by_status(conn, where, params)

            return result
        finally:
            conn.close()

    def _aggregate_by_tool(
        self, conn: sqlite3.Connection, where: str, params: list[object]
    ) -> list[dict]:
        """Break down stats per tool from the comma-separated tool_names column."""
        rows = conn.execute(
            f"SELECT id, tool_names, duration_seconds, total_tokens, total_cost, violations_count FROM sessions {where}",
            params,
        ).fetchall()

        tool_stats: dict[str, dict] = {}
        for row in rows:
            tools = [t.strip() for t in (row["tool_names"] or "").split(",") if t.strip()]
            if not tools:
                tools = ["unknown"]
            for tool in tools:
                if tool not in tool_stats:
                    tool_stats[tool] = {
                        "tool": tool,
                        "sessions": 0,
                        "total_duration": 0.0,
                        "total_tokens": 0,
                        "total_cost": 0.0,
                        "violations_count": 0,
                    }
                tool_stats[tool]["sessions"] += 1
                tool_stats[tool]["total_duration"] += row["duration_seconds"] or 0
                tool_stats[tool]["total_tokens"] += row["total_tokens"] or 0
                tool_stats[tool]["total_cost"] += row["total_cost"] or 0
                tool_stats[tool]["violations_count"] += row["violations_count"] or 0

        return sorted(tool_stats.values(), key=lambda g: g["sessions"], reverse=True)

    def _aggregate_by_status(
        self, conn: sqlite3.Connection, where: str, params: list[object]
    ) -> list[dict]:
        rows = conn.execute(
            f"""\
            SELECT
                status,
                COUNT(*) as sessions,
                COALESCE(SUM(duration_seconds), 0) as total_duration,
                COALESCE(SUM(total_tokens), 0) as total_tokens,
                COALESCE(SUM(total_cost), 0) as total_cost,
                COALESCE(SUM(violations_count), 0) as violations_count
            FROM sessions {where}
            GROUP BY status
            ORDER BY sessions DESC
            """,
            params,
        ).fetchall()
        return [dict(row) for row in rows]

    def rebuild_index(self) -> int:
        """Drop and rebuild the SQLite index from JSON files.

        Returns the number of sessions indexed.
        """
        conn = self._connect()
        try:
            conn.execute("DELETE FROM sessions")
            self._backfill(conn)
            count = conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0]
            conn.commit()
            return count
        finally:
            conn.close()
