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


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.assertEqual(signal.grouping_key, f"reviewer:{self.task.id}:REWORK_REQUIRED")
        
        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, "MALFORMED_OUTPUT")
        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, "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_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")