"""File-backed storage layer for GoalState objects."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING

from mcp_vm_blackbox._helpers import _atomic_write_json, default_blackbox_data_root
from mcp_vm_blackbox.models import GoalState

if TYPE_CHECKING:
    from pathlib import Path

_GITIGNORE_CONTENT = "*\n!.gitignore\n"

_DEFAULT_ROOT = default_blackbox_data_root()


def composite_key(vm_name: str, scenario_name: str) -> str:
    """Build the canonical vm_name/scenario_name composite key."""
    return f"{vm_name}/{scenario_name}"


def run_status_key(ck: str) -> str:
    """Build the session-state key for run status."""
    return f"run_status:{ck}"


def goal_state_key(ck: str) -> str:
    """Build the session-state key for goal state."""
    return f"goal:{ck}"


class GoalStateStore:
    """File-backed storage for GoalState objects.

    Storage layout:
      {root}/goal-states/{vm_name}/{scenario_name}.json        # committed
      {root}/goal-states/{vm_name}/{scenario_name}.draft.json  # in-progress/aborted
    """

    def __init__(self, root: Path) -> None:
        """Initialize the store.

        Args:
            root: Repo-local storage root (typically from :func:`default_blackbox_data_root`).
        """
        self._root = root
        self._states_dir = root / "goal-states"
        self._states_dir.mkdir(parents=True, exist_ok=True)
        gitignore = root / ".gitignore"
        if not gitignore.exists():
            gitignore.write_text(_GITIGNORE_CONTENT, encoding="utf-8")

    def _committed_path(self, vm_name: str, scenario_name: str) -> Path:
        return self._states_dir / vm_name / f"{scenario_name}.json"

    def _draft_path(self, vm_name: str, scenario_name: str) -> Path:
        return self._states_dir / vm_name / f"{scenario_name}.draft.json"

    def save(self, goal: GoalState) -> None:
        """Write committed goal state atomically. Overwrites existing file."""
        path = self._committed_path(goal.vm_name, goal.scenario_name)
        _atomic_write_json(path, goal.model_dump())

    def save_draft(self, goal: GoalState) -> None:
        """Write draft goal state atomically. goal.status must be 'draft'."""
        path = self._draft_path(goal.vm_name, goal.scenario_name)
        _atomic_write_json(path, goal.model_dump())

    def promote_draft(self, vm_name: str, scenario_name: str, draft: GoalState | None = None) -> GoalState:
        """Rename .draft.json to .json and set status='committed'.

        Args:
            vm_name: VM name component of the composite key.
            scenario_name: Scenario name component.
            draft: In-memory draft to promote. When provided, avoids reading the
                draft file from disk (eliminates the read-after-write round-trip).
                When None, the draft file is read from disk.

        Raises:
            FileNotFoundError: if draft file does not exist (when draft=None).
        """
        draft_path = self._draft_path(vm_name, scenario_name)
        if draft is not None:
            data = draft.model_dump()
        else:
            if not draft_path.exists():
                raise FileNotFoundError(f"Draft file not found: {draft_path}")
            data = json.loads(draft_path.read_text(encoding="utf-8"))
        data["status"] = "committed"
        committed_path = self._committed_path(vm_name, scenario_name)
        _atomic_write_json(committed_path, data)
        if draft_path.exists():
            draft_path.unlink()
        return GoalState.model_validate(data)

    def load(self, vm_name: str, scenario_name: str) -> GoalState | None:
        """Load committed goal state. Returns None if file does not exist."""
        path = self._committed_path(vm_name, scenario_name)
        if not path.exists():
            return None
        data = json.loads(path.read_text(encoding="utf-8"))
        return GoalState.model_validate(data)

    def load_draft(self, vm_name: str, scenario_name: str) -> GoalState | None:
        """Load draft goal state. Returns None if draft file does not exist."""
        path = self._draft_path(vm_name, scenario_name)
        if not path.exists():
            return None
        data = json.loads(path.read_text(encoding="utf-8"))
        return GoalState.model_validate(data)

    def load_all(self) -> list[GoalState]:
        """Load all committed goal states from the storage root. Skips drafts."""
        results = []
        for path in self._states_dir.glob("**/*.json"):
            if path.name.endswith(".draft.json"):
                continue
            data = json.loads(path.read_text(encoding="utf-8"))
            results.append(GoalState.model_validate(data))
        return results

    def list_drafts(self) -> list[tuple[str, str]]:
        """Return (vm_name, scenario_name) pairs for all draft files."""
        results: list[tuple[str, str]] = []
        for draft_path in self._states_dir.glob("**/*.draft.json"):
            parts = draft_path.relative_to(self._states_dir).parts
            if len(parts) >= 2:  # noqa: PLR2004
                vm = parts[0]
                scenario = parts[1].replace(".draft.json", "")
                results.append((vm, scenario))
        return results

    def delete(self, vm_name: str, scenario_name: str) -> None:
        """Delete committed file if it exists. No-op if missing."""
        path = self._committed_path(vm_name, scenario_name)
        path.unlink(missing_ok=True)


# Shared singleton — import this instead of constructing GoalStateStore directly.
default_store = GoalStateStore(root=_DEFAULT_ROOT)
