"""Shared state management for CLI agent hook integrations.

Provides a base dataclass, session-id validation, mkdir-based locking,
and state persistence.  Each agent module defines its own HookState
subclass and a ``_state_dir()`` callable, then delegates to these helpers.
"""
from __future__ import annotations

import json
import re
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

_SAFE_SESSION_ID = re.compile(r"^[a-zA-Z0-9_\-]{1,256}$")
_STALE_LOCK_TIMEOUT = 60  # seconds (increased from 30 to avoid racing with slow handle_stop)


def validate_session_id(session_id: str) -> str:
    """Validate session_id to prevent path traversal attacks."""
    if not _SAFE_SESSION_ID.match(session_id):
        raise ValueError(f"Invalid session_id: {session_id!r}")
    return session_id


# ---------------------------------------------------------------------------
# Base state dataclass
# ---------------------------------------------------------------------------


@dataclass
class BaseHookState:
    """Per-session state shared across hook invocations."""

    session_id: str
    cwd: str
    turn_count: int = 0
    current_trace_id: str | None = None
    current_turn_span_id: str | None = None
    current_turn_start: float | None = None
    current_turn_prompt: str = ""
    pending_tool_starts: dict[str, float] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "BaseHookState":
        known = set(cls.__dataclass_fields__)  # type: ignore[attr-defined]
        return cls(**{k: v for k, v in data.items() if k in known})


# ---------------------------------------------------------------------------
# Persistence helpers — all take an explicit *state_dir* so each agent
# module can keep a monkeypatch-friendly ``_state_dir()`` callable.
# ---------------------------------------------------------------------------


def state_path(state_dir: Path, session_id: str) -> Path:
    return state_dir / f"{validate_session_id(session_id)}.json"


def lock_path(state_dir: Path, session_id: str) -> Path:
    return state_dir / f"{validate_session_id(session_id)}.lock"


def load_state_from(
    state_dir: Path, state_cls: type[BaseHookState], session_id: str
) -> BaseHookState | None:
    """Load session state from disk.  Returns None if not found or corrupted."""
    path = state_path(state_dir, session_id)
    if not path.exists():
        return None
    try:
        return state_cls.from_dict(json.loads(path.read_text(encoding="utf-8")))
    except (json.JSONDecodeError, TypeError, KeyError):
        return None


def save_state_to(state_dir: Path, state: BaseHookState) -> None:
    """Persist session state to disk."""
    state_dir.mkdir(parents=True, exist_ok=True)
    state_path(state_dir, state.session_id).write_text(
        json.dumps(state.to_dict(), indent=2), encoding="utf-8"
    )


def delete_state_from(state_dir: Path, session_id: str) -> None:
    state_path(state_dir, session_id).unlink(missing_ok=True)


# ---------------------------------------------------------------------------
# Locking
# ---------------------------------------------------------------------------


def acquire_lock_in(
    state_dir: Path,
    session_id: str,
    max_retries: int = 30,
    retry_delay: float = 0.1,
) -> bool:
    """Acquire a mkdir-based lock.  Returns True on success."""
    state_dir.mkdir(parents=True, exist_ok=True)
    lock = lock_path(state_dir, session_id)
    for _ in range(max_retries):
        try:
            lock.mkdir()
            return True
        except FileExistsError:
            try:
                if time.time() - lock.stat().st_mtime > _STALE_LOCK_TIMEOUT:
                    lock.rmdir()
                    continue
            except OSError:
                pass
            time.sleep(retry_delay)
    return False


def release_lock_in(state_dir: Path, session_id: str) -> None:
    lock = lock_path(state_dir, session_id)
    try:
        lock.rmdir()
    except OSError:
        pass


# ---------------------------------------------------------------------------
# Garbage collection (for agents without SessionEnd, e.g. Codex)
# ---------------------------------------------------------------------------


def gc_stale_in(state_dir: Path, max_age_seconds: int = 86400) -> None:
    """Remove state files and stale lock dirs older than *max_age_seconds*."""
    if not state_dir.exists():
        return
    now = time.time()
    for path in state_dir.glob("*.json"):
        try:
            if now - path.stat().st_mtime > max_age_seconds:
                path.unlink()
        except OSError:
            pass
    for path in state_dir.glob("*.lock"):
        try:
            if path.is_dir() and now - path.stat().st_mtime > max_age_seconds:
                path.rmdir()
        except OSError:
            pass
