"""Tests for the v3 annotation schema, validators, and safety gates."""

from __future__ import annotations

import pytest
from pydantic import ValidationError as PydanticValidationError

from services.intelligence_pipeline_v3.schemas.annotations import (
    AmbiguityType,
    AnnotatedDocument,
    AnnotationMetadata,
    CompanySentimentAnnotation,
    EntityAnnotation,
    EntityType,
    EventClass,
    EvidenceSpanAnnotation,
    RelationAnnotation,
    RelationType,
    SentimentLabel,
)
from services.intelligence_pipeline_v3.schemas.safety import (
    SAFETY_CRITICAL_FIELDS,
    SafetyCriticalField,
    check_safety_gates,
)
from services.intelligence_pipeline_v3.schemas.samples import (
    SAMPLE_BUILDERS,
    build_sample_earnings_beat,
    build_sample_macro_event,
    build_sample_multi_company_competitive,
)
from services.intelligence_pipeline_v3.schemas.validators import (
    validate_annotation,
)

# ---------------------------------------------------------------------------
# Schema model tests
# ---------------------------------------------------------------------------


class TestEvidenceSpan:
    def test_valid_span(self):
        span = EvidenceSpanAnnotation(
            start_char=0, end_char=10, text="Apple Inc."
        )
        assert span.start_char == 0
        assert span.end_char == 10

    def test_end_must_exceed_start(self):
        with pytest.raises(PydanticValidationError):
            EvidenceSpanAnnotation(start_char=10, end_char=5, text="x")

    def test_equal_start_end_rejected(self):
        with pytest.raises(PydanticValidationError):
            EvidenceSpanAnnotation(start_char=5, end_char=5, text="x")

    def test_negative_start_rejected(self):
        with pytest.raises(PydanticValidationError):
            EvidenceSpanAnnotation(start_char=-1, end_char=5, text="hello")


class TestEntityAnnotation:
    def test_requires_evidence(self):
        with pytest.raises(PydanticValidationError):
            EntityAnnotation(
                entity_type=EntityType.COMPANY,
                literal_text="Apple",
                evidence_ids=[],
                confidence=0.9,
            )

    def test_confidence_bounds(self):
        with pytest.raises(PydanticValidationError):
            EntityAnnotation(
                entity_type=EntityType.COMPANY,
                literal_text="Apple",
                evidence_ids=["ev-1"],
                confidence=1.5,
            )


class TestCompanySentiment:
    def test_valid_sentiment(self):
        sent = CompanySentimentAnnotation(
            company_entity_id="ent-1",
            label=SentimentLabel.POSITIVE,
            positive_probability=0.8,
            negative_probability=0.1,
            neutral_probability=0.1,
            evidence_ids=["ev-1"],
            confidence=0.9,
        )
        assert sent.label == SentimentLabel.POSITIVE

    def test_probabilities_must_sum_to_one(self):
        with pytest.raises(PydanticValidationError, match="sum to"):
            CompanySentimentAnnotation(
                company_entity_id="ent-1",
                label=SentimentLabel.POSITIVE,
                positive_probability=0.5,
                negative_probability=0.1,
                neutral_probability=0.1,
                evidence_ids=["ev-1"],
                confidence=0.9,
            )

    def test_allows_small_rounding_error(self):
        # 0.33 + 0.33 + 0.34 = 1.0 exactly, but 0.333+0.333+0.334=1.0 too
        sent = CompanySentimentAnnotation(
            company_entity_id="ent-1",
            label=SentimentLabel.NEUTRAL,
            positive_probability=0.33,
            negative_probability=0.33,
            neutral_probability=0.34,
            evidence_ids=["ev-1"],
            confidence=0.8,
        )
        assert sent.label == SentimentLabel.NEUTRAL


class TestEventAnnotation:
    def test_all_event_classes_defined(self):
        expected = {
            "earnings_beat", "earnings_miss", "guidance_raise", "guidance_cut",
            "ma_announcement", "legal_regulatory", "product_launch", "supply_chain",
            "rating_change", "management_change", "macro_event", "dividend_change",
            "buyback",
        }
        actual = {e.value for e in EventClass}
        assert actual == expected


class TestRelationAnnotation:
    def test_all_relation_types_defined(self):
        expected = {"directly_affects", "inferred_exposure", "competes_with", "supplies"}
        actual = {r.value for r in RelationType}
        assert actual == expected


# ---------------------------------------------------------------------------
# Validator tests
# ---------------------------------------------------------------------------


class TestValidator:
    def test_all_samples_valid(self):
        for builder in SAMPLE_BUILDERS:
            doc = builder()
            result = validate_annotation(doc)
            assert result.valid, f"Sample {doc.document_id} failed: {[e.message for e in result.errors]}"

    def test_detects_invalid_evidence_reference(self):
        doc = build_sample_earnings_beat()
        # Add an entity with a bad evidence reference
        doc.entities.append(
            EntityAnnotation(
                entity_type=EntityType.PERSON,
                literal_text="Tim Cook",
                evidence_ids=["nonexistent-id"],
                confidence=0.9,
            )
        )
        result = validate_annotation(doc)
        assert not result.valid
        assert any("nonexistent-id" in e.message for e in result.errors)

    def test_detects_offset_beyond_text(self):
        source = "Short text."
        doc = AnnotatedDocument(
            document_id="test-doc",
            document_type="article",
            source_text=source,
            metadata=AnnotationMetadata(annotator_id="test"),
            evidence_spans=[
                EvidenceSpanAnnotation(
                    id="ev-bad",
                    start_char=0,
                    end_char=999,
                    text="Short text.",
                )
            ],
        )
        result = validate_annotation(doc)
        assert not result.valid
        assert any("exceeds source_text length" in e.message for e in result.errors)

    def test_detects_text_mismatch(self):
        source = "Apple Inc. beat expectations."
        doc = AnnotatedDocument(
            document_id="test-doc",
            document_type="article",
            source_text=source,
            metadata=AnnotationMetadata(annotator_id="test"),
            evidence_spans=[
                EvidenceSpanAnnotation(
                    id="ev-mismatch",
                    start_char=0,
                    end_char=10,
                    text="Google LLC",  # Doesn't match source
                )
            ],
        )
        result = validate_annotation(doc)
        assert not result.valid
        assert any("does not match" in e.message for e in result.errors)

    def test_warns_on_orphaned_evidence(self):
        source = "Some text here."
        doc = AnnotatedDocument(
            document_id="test-doc",
            document_type="article",
            source_text=source,
            metadata=AnnotationMetadata(annotator_id="test"),
            evidence_spans=[
                EvidenceSpanAnnotation(
                    id="ev-orphan",
                    start_char=0,
                    end_char=4,
                    text="Some",
                )
            ],
        )
        result = validate_annotation(doc)
        assert result.valid  # Warnings don't invalidate
        assert result.warning_count > 0
        assert any("not referenced" in w.message for w in result.warnings)

    def test_detects_invalid_relation_target(self):
        doc = build_sample_multi_company_competitive()
        doc.relations.append(
            RelationAnnotation(
                relation_type=RelationType.SUPPLIES,
                source_id="ent-101",
                target_id="nonexistent-entity",
                evidence_ids=["ev-101"],
                confidence=0.8,
            )
        )
        result = validate_annotation(doc)
        assert not result.valid
        assert any("nonexistent-entity" in e.message for e in result.errors)


# ---------------------------------------------------------------------------
# Safety gate tests
# ---------------------------------------------------------------------------


class TestSafetyGates:
    def test_all_fields_have_thresholds(self):
        for field in SafetyCriticalField:
            assert field in SAFETY_CRITICAL_FIELDS

    def test_passing_metrics(self):
        metrics = {
            SafetyCriticalField.COMPANY_IDENTITY: {"precision": 0.96, "recall": 0.91, "f1": 0.93},
            SafetyCriticalField.EVENT_CLASS: {"macro_f1": 0.87, "per_class_min_f1": 0.72},
            SafetyCriticalField.SENTIMENT_DIRECTION: {"macro_f1": 0.86, "direction_accuracy": 0.91},
            SafetyCriticalField.NUMERIC_FACT_VALUE: {"exact_match": 0.82, "tolerance_match_5pct": 0.93},
            SafetyCriticalField.DIRECT_EFFECT_ATTRIBUTION: {"precision": 0.94, "recall": 0.89},
            SafetyCriticalField.EVIDENCE_SUPPORT: {"support_rate": 0.96, "offset_validity": 0.99},
            SafetyCriticalField.CONFIDENCE_CALIBRATION: {"ece": 0.04, "brier_score": 0.12},
        }
        results = check_safety_gates(metrics)
        assert all(r.passed for r in results), [
            f"{r.field.value}.{r.metric_name}: {r.actual_value} vs {r.required_value}"
            for r in results if not r.passed
        ]

    def test_failing_metrics(self):
        metrics = {
            SafetyCriticalField.COMPANY_IDENTITY: {"precision": 0.80, "recall": 0.70, "f1": 0.75},
        }
        results = check_safety_gates(metrics)
        # All company_identity checks should fail
        company_results = [r for r in results if r.field == SafetyCriticalField.COMPANY_IDENTITY]
        assert all(not r.passed for r in company_results)

    def test_missing_metric_fails(self):
        metrics = {
            SafetyCriticalField.COMPANY_IDENTITY: {"precision": 0.96},  # Missing recall and f1
        }
        results = check_safety_gates(metrics)
        company_results = [r for r in results if r.field == SafetyCriticalField.COMPANY_IDENTITY]
        missing = [r for r in company_results if not r.passed]
        assert len(missing) >= 2  # recall and f1 are missing

    def test_lower_is_better_fields(self):
        """ECE and Brier score are lower-is-better metrics."""
        metrics = {
            SafetyCriticalField.CONFIDENCE_CALIBRATION: {"ece": 0.10, "brier_score": 0.25},
        }
        results = check_safety_gates(metrics)
        cal_results = [r for r in results if r.field == SafetyCriticalField.CONFIDENCE_CALIBRATION]
        assert all(not r.passed for r in cal_results)
        assert all(r.is_lower_better for r in cal_results)


# ---------------------------------------------------------------------------
# Sample annotation tests
# ---------------------------------------------------------------------------


class TestSampleAnnotations:
    def test_earnings_beat_structure(self):
        doc = build_sample_earnings_beat()
        assert doc.document_type == "article"
        assert len(doc.entities) == 1
        assert doc.entities[0].canonical_name == "AAPL"
        assert len(doc.events) == 2
        assert doc.events[0].event_class == EventClass.EARNINGS_BEAT
        assert doc.events[1].event_class == EventClass.DIVIDEND_CHANGE
        assert len(doc.numeric_facts) == 2
        assert len(doc.sentiments) == 1
        assert doc.sentiments[0].label == SentimentLabel.POSITIVE
        assert len(doc.direct_effects) == 1
        assert len(doc.ambiguity_markers) == 0

    def test_multi_company_has_ambiguity(self):
        doc = build_sample_multi_company_competitive()
        assert len(doc.ambiguity_markers) == 1
        assert doc.ambiguity_markers[0].ambiguity_type == AmbiguityType.CONFLICTING_SENTIMENT
        assert len(doc.inferred_exposures) == 1
        assert len(doc.relations) == 1
        assert doc.relations[0].relation_type == RelationType.COMPETES_WITH

    def test_macro_event_no_primary_company(self):
        doc = build_sample_macro_event()
        assert doc.document_type == "macro_event"
        assert doc.events[0].event_class == EventClass.MACRO_EVENT
        assert doc.events[0].primary_company_ids == []
        assert len(doc.sentiments) == 0