from __future__ import annotations

import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

from wov_sdk.models import InvokeRequest, WorkflowDefinition

from app.db import Database
from app.node_manager import NodeManager


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def topological_sort(definition: WorkflowDefinition) -> list[str]:
    nodes = {node.id: node for node in definition.nodes}
    indegree = {node_id: 0 for node_id in nodes}
    dependents: dict[str, list[str]] = {node_id: [] for node_id in nodes}

    for edge in definition.edges:
        if edge.from_node not in nodes or edge.to_node not in nodes:
            raise ValueError(f"unknown edge: {edge.from_node} -> {edge.to_node}")
        indegree[edge.to_node] += 1
        dependents[edge.from_node].append(edge.to_node)

    queue = [node_id for node_id, degree in indegree.items() if degree == 0]
    ordered: list[str] = []
    while queue:
        current = queue.pop(0)
        ordered.append(current)
        for dependent in dependents[current]:
            indegree[dependent] -= 1
            if indegree[dependent] == 0:
                queue.append(dependent)

    if len(ordered) != len(nodes):
        raise ValueError("workflow contains a cycle")
    return ordered


class WorkflowScheduler:
    def __init__(
        self,
        db: Database,
        node_manager: NodeManager,
        storage_dir: Path,
        interval_seconds: float = 1.0,
    ) -> None:
        self.db = db
        self.node_manager = node_manager
        self.storage_dir = storage_dir
        self.interval_seconds = interval_seconds
        self._thread: threading.Thread | None = None
        self._stopping = False

    def start(self) -> None:
        if self._thread is not None:
            return
        self._stopping = False
        self._thread = threading.Thread(
            target=self._loop,
            name="wov-workflow-scheduler",
            daemon=True,
        )
        self._thread.start()

    def stop(self) -> None:
        self._stopping = True
        if self._thread is not None:
            self._thread.join(timeout=5)
            self._thread = None

    def _loop(self) -> None:
        while not self._stopping:
            run = self.db.next_queued_run()
            if run is not None:
                self.execute_run(run["id"])
            else:
                time.sleep(self.interval_seconds)

    def _resolve_ref(
        self,
        ref: str,
        run_input_uri: str | None,
        outputs_by_node: dict[str, dict[str, str]],
    ) -> str | None:
        if ref.startswith("input."):
            return run_input_uri
        node_id, separator, key = ref.partition(".")
        if not separator:
            return None
        return outputs_by_node.get(node_id, {}).get(key)

    def execute_run(self, run_id: str) -> None:
        run = self.db.get_run(run_id)
        if run is None or run["status"] != "QUEUED":
            return

        workflow = self.db.get_workflow(run["workflow_id"])
        if workflow is None:
            self.db.update_run(run_id, status="FAILED", error="workflow not found", updated_at=_now_iso())
            return

        version = self.db.get_workflow_version(run["workflow_id"], run["workflow_version"])
        if version is None:
            self.db.update_run(run_id, status="FAILED", error="workflow version not found", updated_at=_now_iso())
            return

        definition = WorkflowDefinition.from_dict(version["definition"])
        definition.validate()
        ordered = topological_sort(definition)
        outputs_by_node: dict[str, dict[str, str]] = {}
        self.db.update_run(run_id, status="RUNNING", progress=0, updated_at=_now_iso())

        try:
            for index, node_id in enumerate(ordered):
                node_spec = next(item for item in definition.nodes if item.id == node_id)
                self.db.update_run(
                    run_id,
                    current_node_id=node_id,
                    progress=index / len(ordered),
                    updated_at=_now_iso(),
                )
                invoke_inputs: dict[str, str] = {}
                for input_name, ref in node_spec.inputs.items():
                    value = self._resolve_ref(ref, run.get("input_uri"), outputs_by_node)
                    if value is None:
                        raise ValueError(f"missing input {input_name} for node {node_id}")
                    invoke_inputs[input_name] = value

                output_dir = (
                    self.storage_dir
                    / "runs"
                    / run_id
                    / "steps"
                    / node_id
                )
                response = self.node_manager.invoke(
                    node_spec.node_type,
                    InvokeRequest(
                        run_id=run_id,
                        node_instance_id="",
                        inputs=invoke_inputs,
                        params=node_spec.params,
                        output_dir=str(output_dir),
                    ),
                )
                if response.status != "completed":
                    raise RuntimeError(response.error or f"node {node_id} failed")

                outputs_by_node[node_id] = {
                    str(key): str(value) for key, value in response.outputs.items()
                }
                for key, uri in outputs_by_node[node_id].items():
                    artifact = {
                        "run_id": run_id,
                        "node_id": node_id,
                        "name": f"{node_id}.{key}",
                        "uri": uri,
                        "mime_type": self._mime_type(uri),
                        "size": self._file_size(uri),
                    }
                    self.db.create_artifact(artifact)

            for alias, ref in definition.final_outputs.items():
                resolved = self._resolve_ref(ref, run.get("input_uri"), outputs_by_node)
                if resolved is not None:
                    self.db.create_artifact(
                        {
                            "run_id": run_id,
                            "node_id": ref.partition(".")[0],
                            "name": alias,
                            "uri": resolved,
                            "mime_type": self._mime_type(resolved),
                            "size": self._file_size(resolved),
                        }
                    )

            self.db.update_run(
                run_id,
                status="COMPLETED",
                current_node_id=None,
                progress=1.0,
                updated_at=_now_iso(),
            )
        except Exception as exc:  # noqa: BLE001
            self.db.update_run(
                run_id,
                status="FAILED",
                error=str(exc),
                updated_at=_now_iso(),
            )

    @staticmethod
    def _mime_type(uri: str) -> str:
        path = Path(uri)
        suffix = path.suffix.lower()
        return {
            ".srt": "application/x-subrip",
            ".ass": "text/plain",
            ".wav": "audio/wav",
            ".mp4": "video/mp4",
            ".txt": "text/plain",
        }.get(suffix, "application/octet-stream")

    @staticmethod
    def _file_size(uri: str) -> int:
        try:
            return Path(uri).stat().st_size
        except OSError:
            return 0