"""Orchestrator state persistence.

State is saved atomically to orchestrator.json at every phase transition,
enabling crash recovery via ``cleave run --resume``.

Uses JSON (not YAML) for deterministic round-trip serialization. YAML's
implicit type coercion (e.g. ``"true"`` → bool, ``"123"`` → int) causes
data corruption on save/load cycles. JSON preserves types exactly.
"""

from __future__ import annotations

import json
import logging
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path

from cleave.core.file_utils import atomic_write_text
from cleave.orchestrator.config import OrchestratorConfig
from cleave.orchestrator.errors import ResumeError

logger = logging.getLogger(__name__)

# Valid phase transitions (current → allowed next)
PHASE_ORDER = [
    "preflight",
    "planning",
    "planned",
    "assessing",
    "initializing",
    "dispatching",
    "monitoring",
    "harvesting",
    "reunifying",
    "merging",
    "complete",
    "failed",
]


@dataclass
class ChildState:
    """Tracks one child task throughout its lifecycle."""

    child_id: int
    label: str
    depends_on: list[str] = field(default_factory=list)
    status: str = "pending"  # pending | running | completed | failed | needs_decomposition
    branch_name: str = ""
    worktree_path: str | None = None
    pid: int | None = None
    started_at: str | None = None
    completed_at: str | None = None
    exit_code: int | None = None
    retry_count: int = 0
    cost_usd: float | None = None
    error_message: str | None = None

    def to_dict(self) -> dict:
        return {
            "child_id": self.child_id,
            "label": self.label,
            "depends_on": self.depends_on,
            "status": self.status,
            "branch_name": self.branch_name,
            "worktree_path": self.worktree_path,
            "pid": self.pid,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "exit_code": self.exit_code,
            "retry_count": self.retry_count,
            "cost_usd": self.cost_usd,
            "error_message": self.error_message,
        }

    @classmethod
    def from_dict(cls, data: dict) -> ChildState:
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})


@dataclass
class OrchestratorState:
    """Full orchestrator state, persisted to orchestrator.yaml."""

    run_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    phase: str = "preflight"
    directive: str = ""
    root_directive: str = ""
    repo_path: str = ""
    workspace_path: str = ""
    base_branch: str = ""
    config: dict = field(default_factory=dict)
    children: list[ChildState] = field(default_factory=list)
    plan: dict | None = None
    assessment: dict | None = None
    reunification: dict | None = None
    merge_report: dict | None = None
    pre_dirty_head: str = ""
    total_cost_usd: float = 0.0
    consecutive_failures: int = 0
    depth: int = 0
    parent_run_id: str | None = None
    started_at: str = field(
        default_factory=lambda: datetime.now(UTC).isoformat()
    )
    completed_at: str | None = None

    def set_phase(self, phase: str) -> None:
        """Transition to a new phase."""
        if phase not in PHASE_ORDER:
            raise ValueError(f"Unknown phase: {phase}")
        logger.info("Phase transition: %s → %s", self.phase, phase)
        self.phase = phase

    def to_dict(self) -> dict:
        return {
            "run_id": self.run_id,
            "phase": self.phase,
            "directive": self.directive,
            "root_directive": self.root_directive,
            "repo_path": self.repo_path,
            "workspace_path": self.workspace_path,
            "base_branch": self.base_branch,
            "config": self.config,
            "children": [c.to_dict() for c in self.children],
            "plan": self.plan,
            "assessment": self.assessment,
            "reunification": self.reunification,
            "merge_report": self.merge_report,
            "pre_dirty_head": self.pre_dirty_head,
            "total_cost_usd": self.total_cost_usd,
            "consecutive_failures": self.consecutive_failures,
            "depth": self.depth,
            "parent_run_id": self.parent_run_id,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
        }

    @classmethod
    def from_dict(cls, data: dict) -> OrchestratorState:
        # Shallow copy to avoid mutating caller's dict
        data = dict(data)
        children_raw = data.pop("children", [])
        children = [
            ChildState.from_dict(c) if isinstance(c, dict) else c for c in children_raw
        ]
        # Filter to known fields
        known = cls.__dataclass_fields__
        filtered = {k: v for k, v in data.items() if k in known and k != "children"}
        state = cls(**filtered)
        state.children = children
        return state

    @classmethod
    def from_config(cls, config: OrchestratorConfig, depth: int = 0,
                    parent_run_id: str | None = None) -> OrchestratorState:
        """Create initial state from a config."""
        return cls(
            directive=config.directive,
            root_directive=config.root_directive,
            repo_path=str(config.repo_path),
            config=config.to_dict(),
            depth=depth,
            parent_run_id=parent_run_id,
        )


def save_state(state: OrchestratorState, workspace_path: Path) -> Path:
    """Atomically persist state to orchestrator.json.

    Uses JSON for lossless type preservation. Falls back to writing
    orchestrator.yaml for legacy compatibility if needed.
    """
    state_path = workspace_path / "orchestrator.json"
    content = json.dumps(state.to_dict(), indent=2, ensure_ascii=False)
    atomic_write_text(state_path, content)
    logger.debug("State saved to %s (phase=%s)", state_path, state.phase)
    return state_path


def load_state(workspace_path: Path) -> OrchestratorState:
    """Load state from orchestrator.json, raising ResumeError on failure.

    Also checks for legacy orchestrator.yaml for backward compatibility.
    """
    state_path = workspace_path / "orchestrator.json"

    # Fallback to legacy YAML if JSON doesn't exist
    if not state_path.exists():
        yaml_path = workspace_path / "orchestrator.yaml"
        if yaml_path.exists():
            logger.warning(
                "Found legacy orchestrator.yaml (no orchestrator.json). "
                "Loading from YAML — will re-save as JSON."
            )
            try:
                from cleave.core.yaml_utils import parse_yaml_simple

                content = yaml_path.read_text(encoding="utf-8")
                data = parse_yaml_simple(content)
                state = OrchestratorState.from_dict(data)
                # Re-save as JSON for future loads
                save_state(state, workspace_path)
                logger.info("Migrated state to orchestrator.json")
                return state
            except Exception as e:
                raise ResumeError(f"Failed to parse legacy orchestrator.yaml: {e}") from e
        raise ResumeError(f"No orchestrator.json found in {workspace_path}")

    try:
        content = state_path.read_text(encoding="utf-8")
        data = json.loads(content)
        state = OrchestratorState.from_dict(data)
        logger.info("Loaded state: run_id=%s, phase=%s", state.run_id, state.phase)
        return state
    except json.JSONDecodeError as e:
        raise ResumeError(f"Corrupted orchestrator.json: {e}") from e
    except Exception as e:
        raise ResumeError(f"Failed to load orchestrator.json: {e}") from e


def reset_interrupted_children(state: OrchestratorState) -> int:
    """Reset any children stuck in 'running' back to 'pending' for resume.

    Returns the number of children reset.
    """
    count = 0
    for child in state.children:
        if child.status == "running":
            logger.info("Resetting interrupted child %d (%s) to pending", child.child_id, child.label)
            child.status = "pending"
            child.pid = None
            child.started_at = None
            count += 1
    return count
