from __future__ import annotations

from agents.progeny import ProgenyService
from control_plane.agents.models import Agent, AgentVersion, ImprovementCandidate, ProgenyInvestigation, ProgenySignal, PromotionStatus
from control_plane.projects.models import Milestone, Project, ProjectPlan, Task
from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphRun
from graph.task_execution import task_execution_graph_v1


def fixture_context():
    project = Project.objects.create(name="Progeny Investigation", goal="Improve Artifex")
    plan = ProjectPlan.objects.create(project=project, version=1, goal="Improve Artifex")
    milestone = Milestone.objects.create(project=project, plan=plan, key="P1", title="Progeny", goal="Signals")
    task = Task.objects.create(project=project, milestone=milestone, task_type="implementation", goal="Investigate failure")
    agent = Agent.objects.create(name="Investigation Reviewer", role="REVIEWER")
    version = AgentVersion.objects.create(agent=agent, version=1, model="qwen", system_contract="Review", promotion_status=PromotionStatus.CHAMPION)
    spec = task_execution_graph_v1()
    definition = ExecutionGraphDefinition.objects.create(name="investigation_task_execution", graph_type=spec.graph_type)
    graph_version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
    graph_run = GraphRun.objects.create(execution_graph_version=graph_version, project=project, milestone=milestone, task=task, current_node="review")
    review_node = graph_run.node_runs.create(node_id="review", node_type="review", visit_index=1)
    return project, milestone, task, version, graph_version, graph_run, review_node


def test_smart_investigation_distinguishes_real_failure_patterns() -> None:
    project, milestone, task, agent_version, graph_version, graph_run, review_node = fixture_context()
    service = ProgenyService()
    cases = [
        ("tool_policy", "retry", "TASK_RETRY_EXHAUSTED", "Unsupported operation: delete_file", {"classification": "missing_capability"}, "TOOL_POLICY"),
        ("model_output", "model_output", "MODEL_OUTPUT_INVALID", "Provider returned malformed JSON", {"raw_output": "{invalid"}, "MODEL"),
        ("infrastructure", "provider", "PROVIDER_TIMEOUT", "Spark PATH test runner environment failure", {"stderr": "python not found on PATH"}, "INFRASTRUCTURE"),
        ("reviewer", "reviewer", "REJECTED", "Reviewer false rejection around health route", {"findings": [{"type": "false_rejection"}]}, "REVIEWER"),
    ]

    for key, source, category, summary, evidence, expected_target in cases:
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=review_node,
            execution_graph_version=graph_version,
            source=source,
            severity="high",
            failure_category=category,
            summary=summary,
            evidence=evidence,
            grouping_key=key,
            model=agent_version.model,
        )

        investigation = service.create_smart_investigation(signal.grouping_key)

        assert investigation.recommended_target == expected_target
        assert investigation.signals.get() == signal


def test_smart_investigation_can_recommend_graph_node_and_create_graph_candidate() -> None:
    project, milestone, task, agent_version, graph_version, graph_run, review_node = fixture_context()
    service = ProgenyService()
    for index in range(2):
        ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=review_node,
            execution_graph_version=graph_version,
            source="graph_runtime",
            severity="high",
            failure_category="GRAPH_RUNTIME_ERROR",
            summary="Repeated failure concentrated at review node",
            evidence={"node_id": "review", "index": index},
            grouping_key="graph-node-review",
            model=agent_version.model,
        )

    investigation = service.create_smart_investigation("graph-node-review")
    candidate = service.create_improvement_candidate(
        investigation,
        hypothesis="Insert static-analysis node before Reviewer to reduce REWORK_REQUIRED",
    )

    assert investigation.recommended_target == "GRAPH_NODE"
    assert investigation.recommended_route == "Progeny Graph Evolution"
    assert investigation.affected_graph_versions == [str(graph_version.id)]
    assert investigation.affected_nodes == ["review"]
    assert candidate.target_type == "GRAPH_NODE"
    assert candidate.execution_graph_version == graph_version
    assert candidate.hypothesis == "Insert static-analysis node before Reviewer to reduce REWORK_REQUIRED"
    assert ProgenyInvestigation.objects.count() == 1
    assert ImprovementCandidate.objects.count() == 1