"""Session persistence and resume for the Hikaflow CLI.

Saves conversation state (messages + metadata) as JSONL to
.hikaflow/sessions/ so sessions can be resumed after terminal close
or interruption.

Inspired by Codex CLI's resume_picker.rs and Gemini CLI's sessionUtils.ts.
"""

from __future__ import annotations

import json
import logging
import os
import subprocess
import time
import uuid
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)

SESSIONS_DIR_NAME = "sessions"
_HIKAFLOW_DIR = ".hikaflow"


@dataclass
class SessionMeta:
    """Metadata for a saved session."""
    session_id: str
    created_at: float
    updated_at: float
    mode: str
    model: str
    cwd: str
    git_branch: str
    preview: str          # First 120 chars of initial prompt
    turn_count: int = 0
    token_total: int = 0


def _sessions_dir() -> Path:
    """Get sessions directory, creating if needed."""
    d = Path.home() / _HIKAFLOW_DIR / SESSIONS_DIR_NAME
    d.mkdir(parents=True, exist_ok=True)
    return d


def _get_git_branch() -> str:
    """Get current git branch name."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
            capture_output=True, text=True, timeout=3,
        )
        if result.returncode == 0:
            return result.stdout.strip()
    except Exception:
        pass
    return ""


def _serialize_message(msg: Any) -> dict:
    """Serialize a LangChain message to a JSON-safe dict."""
    msg_type = type(msg).__name__
    content = getattr(msg, "content", "")
    data = {"type": msg_type, "content": content}

    # Preserve tool_call_id for ToolMessage
    if hasattr(msg, "tool_call_id"):
        data["tool_call_id"] = msg.tool_call_id

    # Preserve tool_calls for AIMessage
    if hasattr(msg, "tool_calls") and msg.tool_calls:
        data["tool_calls"] = [
            {
                "id": tc.get("id", ""),
                "name": tc.get("name", ""),
                "args": tc.get("args", {}),
            }
            for tc in msg.tool_calls
        ]

    # Preserve additional_kwargs if present (e.g., refusal, thinking)
    if hasattr(msg, "additional_kwargs") and msg.additional_kwargs:
        # Only keep serializable keys
        safe = {}
        for k, v in msg.additional_kwargs.items():
            try:
                json.dumps(v)
                safe[k] = v
            except (TypeError, ValueError):
                pass
        if safe:
            data["additional_kwargs"] = safe

    return data


def _deserialize_message(data: dict) -> Any:
    """Deserialize a dict back to a LangChain message."""
    from langchain_core.messages import (
        AIMessage,
        HumanMessage,
        SystemMessage,
        ToolMessage,
    )

    msg_type = data.get("type", "HumanMessage")
    content = data.get("content", "")
    kwargs = data.get("additional_kwargs", {})

    if msg_type == "SystemMessage":
        return SystemMessage(content=content)
    elif msg_type == "HumanMessage":
        return HumanMessage(content=content)
    elif msg_type == "AIMessage":
        msg = AIMessage(content=content, additional_kwargs=kwargs)
        # Restore tool_calls if present
        if "tool_calls" in data:
            msg.tool_calls = data["tool_calls"]
        return msg
    elif msg_type == "ToolMessage":
        return ToolMessage(
            content=content,
            tool_call_id=data.get("tool_call_id", ""),
        )
    else:
        return HumanMessage(content=content)


class SessionStore:
    """Persist and resume debugging sessions.

    Each session is stored as two files:
    - {session_id}.meta.json — metadata (mode, model, timestamps, preview)
    - {session_id}.messages.jsonl — one message per line
    """

    def __init__(self) -> None:
        self._dir = _sessions_dir()

    def new_session_id(self) -> str:
        """Generate a new session ID."""
        return uuid.uuid4().hex[:12]

    def save_meta(self, meta: SessionMeta) -> None:
        """Save or update session metadata."""
        path = self._dir / f"{meta.session_id}.meta.json"
        try:
            path.write_text(
                json.dumps(asdict(meta), indent=2),
                encoding="utf-8",
            )
        except Exception as e:
            logger.debug("Failed to save session meta: %s", e)

    def save_messages(self, session_id: str, messages: list) -> None:
        """Overwrite the message file with current conversation."""
        path = self._dir / f"{session_id}.messages.jsonl"
        try:
            with open(path, "w", encoding="utf-8") as f:
                for msg in messages:
                    line = json.dumps(_serialize_message(msg), default=str)
                    f.write(line + "\n")
        except Exception as e:
            logger.debug("Failed to save session messages: %s", e)

    def append_message(self, session_id: str, msg: Any) -> None:
        """Append a single message (crash-safe incremental save)."""
        path = self._dir / f"{session_id}.messages.jsonl"
        try:
            with open(path, "a", encoding="utf-8") as f:
                line = json.dumps(_serialize_message(msg), default=str)
                f.write(line + "\n")
        except Exception as e:
            logger.debug("Failed to append session message: %s", e)

    def load_session(self, session_id: str) -> tuple[Optional[SessionMeta], list]:
        """Load a session's metadata and messages.

        Returns (meta, messages). meta is None if session not found.
        """
        meta_path = self._dir / f"{session_id}.meta.json"
        msg_path = self._dir / f"{session_id}.messages.jsonl"

        meta = None
        messages = []

        if meta_path.exists():
            try:
                data = json.loads(meta_path.read_text(encoding="utf-8"))
                meta = SessionMeta(**data)
            except Exception as e:
                logger.debug("Failed to load session meta: %s", e)

        if msg_path.exists():
            try:
                for line in msg_path.read_text(encoding="utf-8").splitlines():
                    line = line.strip()
                    if not line:
                        continue
                    data = json.loads(line)
                    messages.append(_deserialize_message(data))
            except Exception as e:
                logger.debug("Failed to load session messages: %s", e)

        return meta, messages

    def list_sessions(self, limit: int = 20) -> list[SessionMeta]:
        """List recent sessions, newest first."""
        sessions = []
        for meta_file in self._dir.glob("*.meta.json"):
            try:
                data = json.loads(meta_file.read_text(encoding="utf-8"))
                sessions.append(SessionMeta(**data))
            except Exception:
                continue

        sessions.sort(key=lambda s: s.updated_at, reverse=True)
        return sessions[:limit]

    def get_latest(self) -> Optional[SessionMeta]:
        """Get the most recent session."""
        sessions = self.list_sessions(limit=1)
        return sessions[0] if sessions else None

    def delete_session(self, session_id: str) -> None:
        """Delete a session's files."""
        for suffix in (".meta.json", ".messages.jsonl"):
            path = self._dir / f"{session_id}{suffix}"
            if path.exists():
                path.unlink()
