from __future__ import annotations

import json
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from uuid import uuid4


def now_iso() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")


@dataclass(frozen=True)
class CurrentContext:
    context: str
    updated_at: str | None


class ContextStore:
    def __init__(self, state_dir: str | Path = ".clfgui") -> None:
        self.state_dir = Path(state_dir)
        self.context_path = self.state_dir / "current_context.json"
        self.runs_path = self.state_dir / "runs.json"
        self.state_dir.mkdir(parents=True, exist_ok=True)

    def get_context(self) -> CurrentContext:
        data = self._read_json(self.context_path, {"context": "", "updated_at": None})
        return CurrentContext(context=str(data.get("context", "")), updated_at=data.get("updated_at"))

    def set_context(self, context: str) -> CurrentContext:
        current = CurrentContext(context=context, updated_at=now_iso())
        self._write_json(self.context_path, {"context": current.context, "updated_at": current.updated_at})
        return current

    def list_runs(self) -> list[dict[str, Any]]:
        runs = self._read_json(self.runs_path, [])
        return runs if isinstance(runs, list) else []

    def get_run(self, run_id: str) -> dict[str, Any] | None:
        for run in self.list_runs():
            if run.get("id") == run_id:
                return run
        return None

    def append_run(
        self,
        *,
        agent_id: str,
        status: str,
        input_context: str,
        output_context: str | None = None,
        error: str | None = None,
    ) -> dict[str, Any]:
        run = {
            "id": str(uuid4()),
            "agent_id": agent_id,
            "status": status,
            "input_context": input_context,
            "output_context": output_context,
            "error": error,
            "created_at": now_iso(),
        }
        runs = self.list_runs()
        runs.append(run)
        self._write_json(self.runs_path, runs)
        return run

    @staticmethod
    def _read_json(path: Path, default):
        if not path.exists():
            return default
        with path.open("r", encoding="utf-8") as handle:
            return json.load(handle)

    @staticmethod
    def _write_json(path: Path, data) -> None:
        tmp_path = path.with_suffix(f"{path.suffix}.tmp")
        with tmp_path.open("w", encoding="utf-8") as handle:
            json.dump(data, handle, ensure_ascii=False, indent=2)
            handle.write("\n")
        tmp_path.replace(path)