from __future__ import annotations

from typing import Any

from django.utils import timezone

from control_plane.authoring.models import ChapterRevision
from graph.bootstrap import champion_story_authoring_graph_v2
from graph.models import GraphRun, GraphRunStatus


class StoryWorkflowRunner:
    def __init__(self, workflow: object) -> None:
        self.workflow = workflow

    def start(self, revision: ChapterRevision, *, max_revisions: int = 2) -> GraphRun:
        version = champion_story_authoring_graph_v2()
        graph_run = GraphRun.objects.create(
            execution_graph_version=version,
            project=revision.chapter.story.project,
            status=GraphRunStatus.RUNNING,
            started_at=timezone.now(),
            current_node="build_context",
            metadata={
                "revision_id": str(revision.id),
                "initial_revision_id": str(revision.id),
                "current_revision_id": str(revision.id),
            },
        )
        thread_id = f"story:{revision.chapter.story_id}:chapter:{revision.chapter.number}:revision:{revision.id}"
        revision.graph_thread_id = thread_id
        revision.save(update_fields=["graph_thread_id", "updated_at"])
        initial = {
            "story_id": str(revision.chapter.story_id),
            "chapter_id": str(revision.chapter_id),
            "revision_id": str(revision.id),
            "graph_run_id": graph_run.id,
            "thread_id": thread_id,
            "editorial_finding_ids": [],
            "patch_finding_ids": [],
            "patch_attempted": False,
            "patch_status": "not_needed",
            "verification_status": "not_needed",
        }
        return self._invoke(graph_run, initial)

    def resume(self, graph_run_id: int, decision: dict[str, Any]) -> GraphRun:
        from langgraph.types import Command

        graph_run = GraphRun.objects.get(id=graph_run_id)
        if graph_run.status == GraphRunStatus.CANCELLED:
            raise RuntimeError("cancelled story runs cannot be resumed")
        graph_run.status = GraphRunStatus.RUNNING
        graph_run.failure_reason = ""
        graph_run.save(update_fields=["status", "failure_reason", "updated_at"])
        value = None if decision.get("action") == "retry" else Command(resume=decision)
        return self._invoke(graph_run, value)

    def _invoke(self, graph_run: GraphRun, value: object) -> GraphRun:
        thread_id = ChapterRevision.objects.get(
            id=graph_run.metadata["revision_id"]
        ).graph_thread_id
        config = {"configurable": {"thread_id": thread_id}}
        try:
            self.workflow.invoke(value, config=config)
            snapshot = self.workflow.get_state(config)
        except Exception as exc:
            graph_run.status = GraphRunStatus.FAILED
            graph_run.failure_reason = str(exc)[:4000]
            graph_run.completed_at = timezone.now()
            graph_run.save(
                update_fields=["status", "failure_reason", "completed_at", "updated_at"]
            )
            raise
        next_nodes = tuple(snapshot.next or ())
        current_revision_id = str(snapshot.values.get("revision_id") or graph_run.metadata["revision_id"])
        graph_run.metadata = {
            **graph_run.metadata,
            "current_revision_id": current_revision_id,
        }
        if next_nodes:
            graph_run.status = GraphRunStatus.PAUSED
            graph_run.current_node = str(next_nodes[0])
            graph_run.failure_reason = "AWAITING_STORY_APPROVAL"
            graph_run.save(
                update_fields=["status", "current_node", "failure_reason", "metadata", "updated_at"]
            )
        else:
            graph_run.status = GraphRunStatus.COMPLETE
            graph_run.current_node = "complete"
            graph_run.completed_at = timezone.now()
            graph_run.metadata = {**graph_run.metadata, "final_state": dict(snapshot.values)}
            graph_run.save(
                update_fields=["status", "current_node", "completed_at", "metadata", "updated_at"]
            )
        return graph_run