"""
JSON state machine runner for workflow templates.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

try:
    from backend.models.agent import RunEvent, RunEventType, TurnContext, WorkflowRun, RunStatus
    from backend.services.tool_registry import ToolRegistry
except ImportError:
    from models.agent import RunEvent, RunEventType, TurnContext, WorkflowRun, RunStatus
    from services.tool_registry import ToolRegistry


class StateMachineRunner:
    def __init__(
        self,
        definition: Dict[str, Any],
        registry: ToolRegistry,
        *,
        on_event: Optional[Callable[[RunEvent], None]] = None,
    ) -> None:
        self.definition = definition
        self.registry = registry
        self.on_event = on_event
        self.states: Dict[str, Dict[str, Any]] = definition.get("states", {})

    @classmethod
    def from_file(cls, path: Path, registry: ToolRegistry, **kwargs) -> "StateMachineRunner":
        with open(path, "r", encoding="utf-8") as f:
            definition = json.load(f)
        return cls(definition, registry, **kwargs)

    def _emit(self, run: WorkflowRun, event_type: RunEventType, **payload: Any) -> RunEvent:
        event = RunEvent(
            run_id=run.id,
            type=event_type,
            state=run.current_state,
            tool=payload.pop("tool", None),
            payload=payload,
        )
        if self.on_event:
            self.on_event(event)
        return event

    async def run(self, run: WorkflowRun, ctx: TurnContext) -> List[RunEvent]:
        events: List[RunEvent] = []
        original_on_event = self.on_event

        def collect(event: RunEvent) -> None:
            events.append(event)
            if original_on_event:
                original_on_event(event)

        self.on_event = collect

        initial = self.definition.get("initial")
        if not initial:
            raise ValueError("State machine missing 'initial' state")

        current = initial
        run.status = RunStatus.RUNNING

        try:
            while current:
                state_def = self.states.get(current)
                if not state_def:
                    raise ValueError(f"Unknown state: {current}")

                run.current_state = current
                events.append(self._emit(run, RunEventType.STATE_ENTER, state=current))

                tool_name = state_def.get("tool")
                if tool_name:
                    events.append(self._emit(run, RunEventType.TOOL_START, tool=tool_name))
                    await self.registry.execute(tool_name, ctx)
                    events.append(
                        self._emit(
                            run,
                            RunEventType.TOOL_END,
                            tool=tool_name,
                            success=True,
                        )
                    )

                current = state_def.get("next")
                if current == "end" or current is None:
                    break

            run.status = RunStatus.COMPLETED
            run.result_content = ctx.generated_content
            events.append(self._emit(run, RunEventType.COMPLETE))
        except Exception as exc:
            run.status = RunStatus.FAILED
            run.error = str(exc)
            ctx.error = str(exc)
            events.append(self._emit(run, RunEventType.ERROR, message=str(exc)))
            raise
        finally:
            self.on_event = original_on_event

        return events