from __future__ import annotations

import pytest
from django.test import TestCase

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


class ProgenySignalTests(TestCase):
    def setUp(self):
        self.project = Project.objects.create(name="Test Project", goal="Test goal")
        self.plan = ProjectPlan.objects.create(project=self.project, version=1, goal="Test plan goal")
        self.milestone = Milestone.objects.create(project=self.project, plan=self.plan, key="M1", title="Milestone 1", goal="Milestone goal")
        self.task = Task.objects.create(project=self.project, milestone=self.milestone, task_type="CODING", goal="Implement feature")
        self.agent = Agent.objects.create(name="Test Agent", role="CODER")
        self.agent_version = AgentVersion.objects.create(
            agent=self.agent,
            version=1,
            model="test-model",
            system_contract="Test contract",
            promotion_status=PromotionStatus.CHAMPION,
        )
        self.service = ProgenyService()

    def test_create_reviewer_rework_signal(self):
        signal = self.service.create_reviewer_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            status="REWORK_REQUIRED",
            findings=[{"type": "tests_failed", "severity": "high", "message": "Tests failed"}],
            summary="Reviewer requested rework",
        )
        self.assertEqual(signal.source, "reviewer")
        self.assertEqual(signal.failure_category, "REWORK_REQUIRED")
        self.assertEqual(signal.severity, "high")
        self.assertIn("reviewer:REWORK_REQUIRED", signal.grouping_key)
        
        event = Event.objects.get(event_type="PROGENY_SIGNAL_CREATED")
        self.assertEqual(event.payload["signal_id"], str(signal.id))
        self.assertEqual(event.payload["source"], "reviewer")

    def test_create_reviewer_reject_signal(self):
        signal = self.service.create_reviewer_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            status="REJECTED",
            findings=[{"type": "empty_diff", "severity": "high", "message": "No diff"}],
            summary="Reviewer rejected",
        )
        self.assertEqual(signal.source, "reviewer")
        self.assertEqual(signal.failure_category, "REJECTED")
        self.assertEqual(signal.severity, "high")
        
        event = Event.objects.get(event_type="PROGENY_SIGNAL_CREATED")
        self.assertEqual(event.payload["status"], "REJECTED")

    def test_create_judge_fail_signal(self):
        signal = self.service.create_judge_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            result="FAIL",
            evidence=[{"type": "test_status", "status": "FAIL"}],
            summary="Judge failed",
        )
        self.assertEqual(signal.source, "judge")
        self.assertEqual(signal.failure_category, "FAIL")
        self.assertEqual(signal.severity, "high")
        
        event = Event.objects.get(event_type="PROGENY_SIGNAL_CREATED")
        self.assertEqual(event.payload["source"], "judge")
        self.assertEqual(event.payload["result"], "FAIL")

    def test_create_model_output_signal(self):
        signal = self.service.create_model_output_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            error="Invalid JSON",
            raw_output="{invalid json}",
        )
        self.assertEqual(signal.source, "model_output")
        self.assertEqual(signal.failure_category, "MODEL_OUTPUT_INVALID")
        self.assertEqual(signal.severity, "high")
        self.assertIn("Invalid JSON", signal.summary)
        
        event = Event.objects.get(event_type="PROGENY_SIGNAL_CREATED")
        self.assertEqual(event.payload["source"], "model_output")

    def test_create_retry_exhausted_signal(self):
        signal = self.service.create_retry_exhausted_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            attempts=3,
        )
        self.assertEqual(signal.source, "retry")
        self.assertEqual(signal.failure_category, "TASK_RETRY_EXHAUSTED")
        self.assertEqual(signal.severity, "critical")
        self.assertIn("exhausted 3 retries", signal.summary)
        
        event = Event.objects.get(event_type="PROGENY_SIGNAL_CREATED")
        self.assertEqual(event.payload["source"], "retry")

    def test_grouping_key_deterministic(self):
        signal1 = self.service.create_reviewer_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            status="REWORK_REQUIRED",
            findings=[],
            summary="Rework 1",
        )
        signal2 = self.service.create_reviewer_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            status="REWORK_REQUIRED",
            findings=[],
            summary="Rework 2",
        )
        self.assertEqual(signal1.grouping_key, signal2.grouping_key)

    def test_signal_records_graph_lineage_when_available(self):
        spec = task_execution_graph_v1()
        definition = ExecutionGraphDefinition.objects.create(name="progeny_lineage", graph_type=spec.graph_type)
        version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
        graph_run = GraphRun.objects.create(execution_graph_version=version, project=self.project, milestone=self.milestone, task=self.task, current_node="review")
        node_run = graph_run.node_runs.create(node_id="review", node_type="review", visit_index=2)

        signal = self.service.create_reviewer_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            status="REWORK_REQUIRED",
            findings=[{"type": "missing_health_route"}],
            summary="Reviewer requested rework",
            graph_run=graph_run,
            graph_node_run=node_run,
        )

        self.assertEqual(signal.graph_run, graph_run)
        self.assertEqual(signal.execution_graph_version, version)
        self.assertEqual(signal.graph_node_run, node_run)
        self.assertEqual(signal.evidence["node_id"], "review")
        self.assertEqual(signal.evidence["node_type"], "review")
        self.assertEqual(signal.evidence["visit_index"], 2)

    def test_inbox_filters_unresolved_signals_by_graph_dimensions(self):
        spec = task_execution_graph_v1()
        definition = ExecutionGraphDefinition.objects.create(name="progeny_inbox", graph_type=spec.graph_type)
        version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
        graph_run = GraphRun.objects.create(execution_graph_version=version, project=self.project, milestone=self.milestone, task=self.task, current_node="judge")
        node_run = graph_run.node_runs.create(node_id="judge", node_type="judge", visit_index=1)
        signal = self.service.create_judge_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            result="FAIL",
            evidence=[{"type": "test_status", "status": "FAIL"}],
            summary="Judge failed",
            graph_run=graph_run,
            graph_node_run=node_run,
        )

        matches = list(
            self.service.query_inbox(
                source="judge",
                status="OPEN",
                project=self.project,
                agent=self.agent,
                agent_version=self.agent_version,
                model="test-model",
                execution_graph=definition,
                execution_graph_version=version,
                graph_node="judge",
                severity="high",
                failure_category="FAIL",
            )
        )

        self.assertEqual(matches, [signal])

    def test_group_unresolved_signals_exposes_impact_summary(self):
        spec = task_execution_graph_v1()
        definition = ExecutionGraphDefinition.objects.create(name="progeny_grouping", graph_type=spec.graph_type)
        version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
        graph_run = GraphRun.objects.create(execution_graph_version=version, project=self.project, milestone=self.milestone, task=self.task, current_node="review")
        node_run = graph_run.node_runs.create(node_id="review", node_type="review", visit_index=1)
        for index in range(2):
            self.service.create_reviewer_signal(
                project=self.project,
                task=self.task,
                milestone=self.milestone,
                agent_version=self.agent_version,
                status="REWORK_REQUIRED",
                findings=[{"type": "missing_health_route", "index": index}],
                summary="Reviewer requested rework",
                graph_run=graph_run,
                graph_node_run=node_run,
            )

        groups = self.service.group_unresolved_signals(source="reviewer", failure_category="REWORK_REQUIRED")

        self.assertEqual(len(groups), 1)
        self.assertEqual(groups[0].occurrence_count, 2)
        self.assertEqual(groups[0].affected_projects, [str(self.project.id)])
        self.assertEqual(groups[0].affected_agents, [str(self.agent_version.id)])
        self.assertEqual(groups[0].affected_graph_versions, [str(version.id)])
        self.assertEqual(groups[0].affected_graph_nodes, ["review"])

    def test_event_lineage(self):
        self.service.create_reviewer_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            status="REWORK_REQUIRED",
            findings=[],
            summary="Rework",
        )
        self.service.create_judge_signal(
            project=self.project,
            task=self.task,
            milestone=self.milestone,
            agent_version=self.agent_version,
            result="FAIL",
            evidence=[],
            summary="Fail",
        )
        
        events = Event.objects.filter(event_type="PROGENY_SIGNAL_CREATED").order_by("created_at")
        self.assertEqual(events.count(), 2)
        self.assertEqual(events[0].payload["source"], "reviewer")
        self.assertEqual(events[1].payload["source"], "judge")