"""Per-conversation dedupe state, persisted to
derived/student-chat/sessions/{conv_id}.json.

The dedupe set lets the orchestrator skip re-retrieving a doc that was
already pulled earlier in the same conversation. Persisted so a Ctrl-C
on the Gradio app doesn't lose continuity for in-flight chats.
"""
from __future__ import annotations
import json
import logging
import os
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from scripts.lib.runtime_config import SESSIONS_DIR

if sys.platform == "win32":
    import msvcrt

    def _try_lock(fileobj) -> bool:
        try:
            msvcrt.locking(fileobj.fileno(), msvcrt.LK_NBLCK, 1)
            return True
        except OSError:
            return False

    def _unlock(fileobj) -> None:
        try:
            fileobj.seek(0)
            msvcrt.locking(fileobj.fileno(), msvcrt.LK_UNLCK, 1)
        except OSError:
            pass
else:
    import fcntl

    def _try_lock(fileobj) -> bool:
        try:
            fcntl.flock(fileobj.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
            return True
        except OSError:
            return False

    def _unlock(fileobj) -> None:
        try:
            fcntl.flock(fileobj.fileno(), fcntl.LOCK_UN)
        except OSError:
            pass


logger = logging.getLogger(__name__)

_SESSIONS_DIR = SESSIONS_DIR


@dataclass
class SessionState:
    conversation_id: str
    seen_doc_ids: set[str] = field(default_factory=set)
    created_at: str = ""
    updated_at: str = ""


def _now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _path_for(conv_id: str, sessions_dir: Path) -> Path:
    return sessions_dir / f"{conv_id}.json"


def load_or_create(conv_id: str, *, sessions_dir: Path = _SESSIONS_DIR) -> SessionState:
    """Load an existing session or return a fresh one. Never raises."""
    path = _path_for(conv_id, sessions_dir)
    if path.is_file():
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            return SessionState(
                conversation_id=data.get("conversation_id", conv_id),
                seen_doc_ids=set(data.get("seen_doc_ids", [])),
                created_at=data.get("created_at", _now_iso()),
                updated_at=data.get("updated_at", _now_iso()),
            )
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(
                "session_state.load_or_create: failed to read %s (%s) — starting fresh",
                path, e,
            )
    now = _now_iso()
    return SessionState(conversation_id=conv_id, created_at=now, updated_at=now)


def record_retrieved(state: SessionState, doc_ids: list[str]) -> None:
    """Add doc_ids to the dedupe set. In-place."""
    state.seen_doc_ids.update(doc_ids)
    state.updated_at = _now_iso()


def persist(state: SessionState, *, sessions_dir: Path = _SESSIONS_DIR) -> None:
    """Atomically write the state to disk via stage-and-rename, with a
    short non-blocking lock to guard against concurrent writes for the
    same conv_id. If lock acquisition fails, log a warning and skip the
    write — the next turn will re-persist."""
    sessions_dir.mkdir(parents=True, exist_ok=True)
    path = _path_for(state.conversation_id, sessions_dir)
    payload = {
        "conversation_id": state.conversation_id,
        "seen_doc_ids": sorted(state.seen_doc_ids),
        "created_at": state.created_at,
        "updated_at": state.updated_at,
    }
    staging = path.with_name(path.name + ".tmp")
    try:
        with open(staging, "w", encoding="utf-8") as f:
            if not _try_lock(f):
                logger.warning(
                    "session_state.persist: lock contention on %s — skipping write",
                    staging,
                )
                # Drop the staging file the with-block would otherwise leave.
                # Use missing_ok in case f.close() races with another process.
                try:
                    f.close()
                    staging.unlink(missing_ok=True)
                except OSError:
                    pass
                return
            try:
                f.write(json.dumps(payload, ensure_ascii=False))
                f.flush()
                os.fsync(f.fileno())
            finally:
                _unlock(f)
        os.replace(staging, path)
    except BaseException:
        if staging.exists():
            try:
                os.unlink(staging)
            except OSError:
                pass
        raise