"""Small JSON state store for the AI Cockpit."""

from __future__ import annotations

import copy
import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo


ROOT = Path(__file__).resolve().parent
DATA_DIR = ROOT / "data"
DEMO_STATE_PATH = DATA_DIR / "demo_state.json"
STATE_PATH = DATA_DIR / "state.json"
EVENTS_PATH = DATA_DIR / "events.jsonl"

REQUIRED_STATE_KEYS = {
    "active_focus",
    "current_work",
    "sidecars",
    "next_allowed_step",
    "todos",
    "open_threads",
}


class StateStore:
    def __init__(
        self,
        state_path: Path = STATE_PATH,
        demo_state_path: Path = DEMO_STATE_PATH,
        events_path: Path = EVENTS_PATH,
    ) -> None:
        self.state_path = state_path
        self.demo_state_path = demo_state_path
        self.events_path = events_path

    def load_state(self) -> dict[str, Any]:
        path = self.state_path if self.state_path.exists() else self.demo_state_path
        with path.open("r", encoding="utf-8") as handle:
            state = json.load(handle)
        validate_state(state, path)
        return copy.deepcopy(state)

    def save_state(self, state: dict[str, Any]) -> None:
        validate_state(state, self.state_path)
        self.state_path.parent.mkdir(parents=True, exist_ok=True)
        with self.state_path.open("w", encoding="utf-8") as handle:
            json.dump(state, handle, ensure_ascii=False, indent=2)
            handle.write("\n")

    def append_event(self, event: dict[str, Any]) -> dict[str, Any]:
        event = copy.deepcopy(event)
        event.setdefault("id", str(uuid.uuid4()))
        event.setdefault("created_at", datetime.now(ZoneInfo("Europe/Berlin")).isoformat(timespec="seconds"))
        event.setdefault("payload", {})

        if not isinstance(event.get("type"), str) or not event["type"]:
            raise ValueError("Event needs a non-empty string 'type'.")
        if not isinstance(event.get("payload"), dict):
            raise ValueError("Event 'payload' must be an object.")

        self.events_path.parent.mkdir(parents=True, exist_ok=True)
        with self.events_path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(event, ensure_ascii=False, sort_keys=True))
            handle.write("\n")
        return event


def validate_state(state: Any, path: Path) -> None:
    if not isinstance(state, dict):
        raise ValueError(f"{path} must contain a JSON object.")
    missing = sorted(REQUIRED_STATE_KEYS - set(state))
    if missing:
        raise ValueError(f"{path} is missing required state keys: {', '.join(missing)}")
    if not isinstance(state["sidecars"], list) or len(state["sidecars"]) != 5:
        raise ValueError(f"{path} must define exactly five sidecars.")
    lead_count = sum(1 for sidecar in state["sidecars"] if sidecar.get("is_lead") is True)
    if lead_count != 1:
        raise ValueError(f"{path} must mark exactly one sidecar as is_lead=true.")


_DEFAULT_STORE = StateStore()


def load_state() -> dict[str, Any]:
    return _DEFAULT_STORE.load_state()


def save_state(state: dict[str, Any]) -> None:
    _DEFAULT_STORE.save_state(state)


def append_event(event: dict[str, Any]) -> dict[str, Any]:
    return _DEFAULT_STORE.append_event(event)