from __future__ import annotations

import json
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from niannian.brain.executive.task_graph import TaskGraph


@dataclass
class TrackRecord:
    track_id: str
    label: str
    regime: str = "companion"
    task: TaskGraph | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "track_id": self.track_id,
            "label": self.label,
            "regime": self.regime,
            "task": self.task.to_dict() if self.task else None,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> TrackRecord:
        task_raw = data.get("task")
        task = TaskGraph.from_dict(task_raw) if isinstance(task_raw, dict) else None
        return cls(
            track_id=str(data.get("track_id") or ""),
            label=str(data.get("label") or ""),
            regime=str(data.get("regime") or "companion"),
            task=task,
        )


class TaskStateStore:
    """小谋 · 跨轮任务状态（state/tasks.json）。"""

    def __init__(self, home: Path) -> None:
        self._path = home / "state" / "tasks.json"
        self._path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        self._tracks: dict[str, TrackRecord] = {}
        self._load()

    def register_track(self, track_id: str, label: str, *, regime: str = "companion") -> TrackRecord:
        with self._lock:
            rec = self._tracks.get(track_id)
            if rec is None:
                rec = TrackRecord(track_id=track_id, label=label, regime=regime)
                self._tracks[track_id] = rec
            else:
                rec.label = label
                rec.regime = regime
            self._persist()
            return rec

    def set_active_task(self, track_id: str, graph: TaskGraph | None) -> None:
        with self._lock:
            rec = self._tracks.setdefault(track_id, TrackRecord(track_id=track_id, label=track_id))
            rec.task = graph
            self._persist()

    def get_active_graph(self, track_id: str) -> TaskGraph | None:
        with self._lock:
            rec = self._tracks.get(track_id)
            return rec.task if rec else None

    def to_executive_api(self) -> dict[str, Any]:
        with self._lock:
            tracks = [r.to_dict() for r in self._tracks.values()]
            active_tasks: list[dict[str, Any]] = []
            for rec in self._tracks.values():
                if rec.task is None:
                    continue
                done, total = rec.task.progress()
                active_tasks.append(
                    {
                        "track_id": rec.track_id,
                        "label": rec.label,
                        "regime": rec.regime,
                        "task_id": rec.task.task_id,
                        "goal": rec.task.goal,
                        "progress": {"done": done, "total": total},
                        "completed": rec.task.is_complete(),
                        "steps": [s.to_dict() for s in rec.task.steps],
                    }
                )
            return {
                "tasks_file": "state/tasks.json",
                "tracks": tracks,
                "active_tasks": active_tasks,
            }

    def get_public_summary(self) -> list[dict[str, Any]]:
        with self._lock:
            out: list[dict[str, Any]] = []
            for rec in self._tracks.values():
                if rec.task is None:
                    continue
                done, total = rec.task.progress()
                running = next((s.label for s in rec.task.steps if s.status == "running"), None)
                pending = next((s.label for s in rec.task.steps if s.status == "pending"), None)
                phase = running or pending or "完成"
                out.append(
                    {
                        "track_id": rec.track_id,
                        "label": rec.label,
                        "goal": rec.task.goal,
                        "phase": phase,
                        "progress": f"{done}/{total}",
                    }
                )
            return out

    def format_task_block(self) -> str:
        lines: list[str] = []
        for item in self.get_public_summary():
            lines.append(
                f"- [{item['label']}] {item['goal']} · {item['phase']} ({item['progress']})"
            )
        return "\n".join(lines)

    def health_summary(self) -> dict[str, Any]:
        with self._lock:
            active = sum(1 for r in self._tracks.values() if r.task is not None)
            return {
                "tasks_file": "state/tasks.json",
                "active_tracks": active,
                "track_count": len(self._tracks),
            }

    def _load(self) -> None:
        if not self._path.is_file():
            return
        with self._path.open(encoding="utf-8") as f:
            raw = json.load(f)
        tracks = raw.get("tracks") if isinstance(raw, dict) else raw
        if not isinstance(tracks, list):
            return
        for item in tracks:
            if isinstance(item, dict) and item.get("track_id"):
                rec = TrackRecord.from_dict(item)
                self._tracks[rec.track_id] = rec

    def _persist(self) -> None:
        payload = {"tracks": [r.to_dict() for r in self._tracks.values()]}
        tmp = self._path.with_suffix(".json.tmp")
        tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        tmp.replace(self._path)
