"""Session state persistence across hook invocations."""
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}$")


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


def _state_dir() -> Path:
    return Path.home() / ".staso" / "claude-code"


def _state_path(session_id: str) -> Path:
    return _state_dir() / f"{_validate_session_id(session_id)}.json"


def _lock_path(session_id: str) -> Path:
    return _state_dir() / f"{_validate_session_id(session_id)}.lock"


@dataclass
class HookState:
    """Per-session state shared across Claude Code hook invocations."""

    session_id: str
    cwd: str
    turn_count: int = 0
    # Current turn (reset after each Stop)
    current_trace_id: str | None = None
    current_turn_span_id: str | None = None
    current_turn_start: float | None = None
    current_turn_prompt: str = ""
    last_transcript_line: int = 0
    pending_tool_starts: dict[str, float] = field(default_factory=dict)
    pending_subagent_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]) -> HookState:
        known = set(cls.__dataclass_fields__)  # type: ignore[attr-defined]
        return cls(**{k: v for k, v in data.items() if k in known})


def load_state(session_id: str) -> HookState | None:
    """Load session state from disk. Returns None if not found or corrupted."""
    path = _state_path(session_id)
    if not path.exists():
        return None
    try:
        return HookState.from_dict(json.loads(path.read_text(encoding="utf-8")))
    except (json.JSONDecodeError, TypeError, KeyError):
        return None


def save_state(state: HookState) -> None:
    """Persist session state to disk."""
    _state_dir().mkdir(parents=True, exist_ok=True)
    _state_path(state.session_id).write_text(
        json.dumps(state.to_dict(), indent=2), encoding="utf-8"
    )


def delete_state(session_id: str) -> None:
    _state_path(session_id).unlink(missing_ok=True)


def acquire_lock(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(session_id)
    for _ in range(max_retries):
        try:
            lock.mkdir()
            return True
        except FileExistsError:
            # Break stale locks older than 30s
            try:
                if time.time() - lock.stat().st_mtime > 30:
                    lock.rmdir()
                    continue
            except OSError:
                pass
            time.sleep(retry_delay)
    return False


def release_lock(session_id: str) -> None:
    lock = _lock_path(session_id)
    try:
        lock.rmdir()
    except OSError:
        pass
