from __future__ import annotations

import json
from copy import deepcopy
from pathlib import Path
from typing import Any

import pytest

from agent_framework.workflows import FileWorkflowRepository, WorkflowActionRegistry, WorkflowRuntime
from agent_framework.workflows.models import WorkflowDefinition

ROOT = Path(__file__).resolve().parents[2]
CASES = ROOT / "tests" / "migration" / "data" / "workflow_yaml_cases"


def _load_cases():
    return [json.loads(p.read_text(encoding="utf-8")) for p in sorted(CASES.glob("*.json"))]


def _contains_subset(expected: Any, observed: Any) -> bool:
    if isinstance(expected, dict):
        return isinstance(observed, dict) and all(
            key in observed and _contains_subset(value, observed[key]) for key, value in expected.items()
        )
    if isinstance(expected, list):
        return isinstance(observed, list) and len(expected) <= len(observed) and all(
            _contains_subset(value, observed[index]) for index, value in enumerate(expected)
        )
    return expected == observed


class _DefinitionRepository:
    def __init__(self, definition):
        self.definition = definition

    def get_active(self, name):
        assert name == self.definition.name
        return self.definition

    def get_version(self, name, version):
        assert (name, version) == (self.definition.name, self.definition.version)
        return self.definition


def _normalized_expected_nodes(case):
    nodes = list((case.get("expect") or {}).get("trace_nodes") or [])
    # Runtime antigo repetia a action imediatamente anterior ao pause. O runtime
    # novo deliberadamente não a repete para impedir side effects duplicados.
    normalized = []
    for node in nodes:
        if normalized and normalized[-1] == node:
            continue
        normalized.append(node)
    return normalized


@pytest.mark.asyncio
@pytest.mark.parametrize("case", _load_cases(), ids=lambda c: c["id"])
async def test_original_workflow_case_on_framework_runtime(case):
    original = FileWorkflowRepository(ROOT / "workflows").get_version(case["workflow"], int(case["version"]))
    raw = original.model_dump(by_alias=True)
    registry = WorkflowActionRegistry()
    calls = []

    # No teste cada node recebe uma action exclusiva. Isso permite reproduzir
    # exatamente os outputs configurados no caso JSON, mesmo quando dois nodes
    # compartilham a mesma action de produção.
    for node in raw["nodes"]:
        node_id = node["id"]
        real_action = node["action"]
        test_action = f"__case__{node_id}"
        node["action"] = test_action

        async def handler(params, state, *, _node=node_id, _action=real_action):
            cfgs = dict(case.get("actions") or {})
            configured = cfgs.get(_node) if isinstance(cfgs.get(_node), dict) else cfgs.get(_action)
            configured = configured if isinstance(configured, dict) else {}
            calls.append({"node_id": _node, "action": _action, "params": deepcopy(params)})
            if configured.get("success") is False:
                raise RuntimeError(str(configured.get("error") or "falha configurada"))
            return deepcopy(configured.get("output") or {"success": True})

        registry.register(test_action, handler)

    definition = WorkflowDefinition.model_validate(raw)
    runtime = WorkflowRuntime(_DefinitionRepository(definition), actions=registry, checkpointer=None, allow_deterministic_fallback=True)
    first = await runtime.arun(definition.name, deepcopy(case.get("input") or {}), version=definition.version)

    start_expect = case.get("expect_start")
    if isinstance(start_expect, dict):
        if start_expect.get("status") == "WAITING_INPUT":
            assert first.status == "PAUSED"
            if "paused_at" in start_expect:
                assert first.pause.get("node") == start_expect["paused_at"]
            expected_input = first.pause.get("expected_input") or {}
            if "expected_input_key" in start_expect:
                assert expected_input.get("key") == start_expect["expected_input_key"]
            if "allowed_values" in start_expect:
                assert expected_input.get("allowed_values") == start_expect["allowed_values"]

    final = first
    resume = case.get("resume")
    if isinstance(resume, dict):
        final = await runtime.aresume(definition.name, first.execution_id, deepcopy(resume), version=definition.version)

    expect = dict(case.get("expect") or {})
    assert final.status == expect.get("status")
    if "last_node" in expect:
        assert final.state.get("current_node") == expect["last_node"]

    # Data do runtime antigo era a união dos outputs. ``vars`` representa essa
    # informação no runtime novo e é a fonte para referências $.vars.* do YAML.
    flattened = {}
    for value in (final.state.get("vars") or {}).values():
        if isinstance(value, dict):
            flattened.update(value)
    if "data_subset" in expect:
        assert _contains_subset(expect["data_subset"], flattened)

    observed_nodes = [item.get("node") for item in final.trace if item.get("action") != "pause_resume"]
    if "trace_nodes" in expect:
        active_node_ids = {node.id for node in definition.nodes}
        expected_nodes = [node for node in _normalized_expected_nodes(case) if node in active_node_ids]
        assert observed_nodes == expected_nodes

    for expected in expect.get("calls", []):
        found = [c for c in calls if c["node_id"] == expected.get("node_id") and ("action" not in expected or c["action"] == expected["action"])]
        assert found, expected
        subset = expected.get("params_subset")
        if subset is not None:
            assert any(_contains_subset(subset, c["params"]) for c in found)