from __future__ import annotations
import json
import os
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Literal

_DEFAULT_BASE = Path(__file__).parent.parent / "sessions"

_SEED: dict = {
    "current_step": 0,
    "route": {},
    "selected_ids": [],
    "indicator_defs": {},
    "bind_results": {},
    "artifacts": {},
    "trace": [],
}


def new_session(
    input_mode: Literal["message", "chips"],
    question: str = "",
    base_dir: Path | None = None,
) -> str:
    base = Path(base_dir) if base_dir else _DEFAULT_BASE
    session_id = str(uuid.uuid4())
    session_dir = base / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    state = {
        "session_id": session_id,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "input_mode": input_mode,
        "question": question,
        **_SEED,
    }
    _write_json(session_dir / "state.json", state)
    return session_id


def write_state(
    session_id: str,
    patch: dict,
    base_dir: Path | None = None,
) -> None:
    base = Path(base_dir) if base_dir else _DEFAULT_BASE
    state_path = base / session_id / "state.json"
    state = json.loads(state_path.read_text())
    for key, value in patch.items():
        if key == "bind_results" and isinstance(value, dict):
            state.setdefault("bind_results", {}).update(value)
        else:
            state[key] = value
    _write_json(state_path, state)


def load_state(session_id: str, base_dir: Path | None = None) -> dict:
    base = Path(base_dir) if base_dir else _DEFAULT_BASE
    return json.loads((base / session_id / "state.json").read_text())


def _write_json(path: Path, data: dict) -> None:
    tmp = path.with_suffix(".json.tmp")
    tmp.write_text(json.dumps(data, indent=2, ensure_ascii=False))
    os.replace(tmp, path)