import threading

from audita.core.chunking import IndexedSegment, TranscriptSection
from audita.core.config import AuditaConfig, ConfigOverrides
from audita.core.errors import AuditaLLMError
from audita.core.schemas import parse_glossary_yaml, parse_transcript_json
from audita.framework.models import CorrectionProposal, ModuleContext, ModuleRunSpec
from audita.framework.runner import PipelineRunner
from audita.validators import (
    MeaningReversalValidator,
    ProposalConfidenceValidator,
    ProtectedGlossaryTermsValidator,
    SpokenFormPlausibilityValidator,
)
from audita.validators.base import ValidationContext, ValidationDecision, ValidationResult


class RecordingValidator:
    execution_kind = "deterministic"

    def __init__(self, name, recorder, approve=True):
        self.name = name
        self._recorder = recorder
        self._approve = approve

    def validate(self, context: ValidationContext) -> ValidationResult:
        self._recorder.append((self.name, [proposal.corrected_text for proposal in context.proposals]))
        return ValidationResult(
            validator_name=self.name,
            execution_kind=self.execution_kind,
            decisions=[
                ValidationDecision(
                    proposal_index=proposal.proposal_index,
                    approved=self._approve,
                    reason=None if self._approve else f"{self.name} rejected proposal",
                )
                for proposal in context.proposals
            ],
        )


class RecordingLLMValidator(RecordingValidator):
    execution_kind = "llm"


class FakeStructuredLLMClient:
    def __init__(self, responses):
        self._responses = responses
        self._lock = threading.Lock()
        self.calls = []

    def run_structured(self, *, stage_name, messages, response_model, config):
        with self._lock:
            self.calls.append(
                {
                    "stage_name": stage_name,
                    "messages": list(messages),
                    "response_model": response_model,
                }
            )
            payload = _pop_llm_response(self._responses, stage_name)
        return response_model.model_validate(payload)


def _pop_llm_response(responses, stage_name):
    if isinstance(responses, dict):
        if stage_name not in responses:
            raise AuditaLLMError(f"FakeStructuredLLMClient received unexpected stage_name: {stage_name}")
        payloads = responses[stage_name]
        if isinstance(payloads, list):
            if not payloads:
                raise AuditaLLMError(f"FakeStructuredLLMClient received too many calls for stage_name: {stage_name}")
            return payloads.pop(0)
        payload = payloads
        del responses[stage_name]
        return payload
    if not responses:
        raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
    return responses.pop(0)


class TrackingStructuredLLMClient:
    def __init__(self, responses, barrier=None):
        self._responses = responses
        self._barrier = barrier
        self._lock = threading.Lock()
        self.calls = []
        self.in_flight = 0
        self.max_in_flight = 0

    def run_structured(self, *, stage_name, messages, response_model, config):
        if self._barrier is not None:
            self._barrier.wait()
        with self._lock:
            self.in_flight += 1
            self.max_in_flight = max(self.max_in_flight, self.in_flight)
            self.calls.append(
                {
                    "stage_name": stage_name,
                    "messages": list(messages),
                    "response_model": response_model,
                }
            )
            payload = _pop_llm_response(self._responses, stage_name)
        try:
            return response_model.model_validate(payload)
        finally:
            with self._lock:
                self.in_flight -= 1


class RecordingModule:
    replacement_policy = "require_unique"

    def __init__(self, module_key, proposals, validators, recorder):
        self.module_key = module_key
        self._proposals = proposals
        self._validators = validators
        self._recorder = recorder

    def validators(self):
        return list(self._validators)

    def propose(self, transcript_section, context: ModuleContext):
        self._recorder.append(("propose", [item.segment.text for item in transcript_section.segments]))
        return list(self._proposals)


class ConcurrentRecordingModule:
    replacement_policy = "require_unique"

    def __init__(self, recorder, barrier):
        self.module_key = "concurrent"
        self._recorder = recorder
        self._barrier = barrier

    def validators(self):
        return []

    def propose(self, transcript_section, context: ModuleContext):
        texts = [item.segment.text for item in transcript_section.segments]
        self._recorder.append(("start", transcript_section.section_index, texts))
        self._barrier.wait()
        segment = transcript_section.segments[0].segment
        return [
            CorrectionProposal(
                proposal_index=0,
                module_instance=context.run_spec.instance_name,
                module_key=context.run_spec.module_key,
                id=segment.id,
                original_text=segment.text.rstrip("."),
                corrected_text=f"{segment.text.rstrip('.')} revised",
                confidence=0.9,
            )
        ]


def test_pipeline_runner_applies_modules_sequentially(tmp_path):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Alpha"
            category: noun
            summary: "Alpha."
        """
    )
    seen = []
    first = RecordingModule(
        "first",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="first",
                module_key="first",
                id=1,
                original_text="Alpha",
                corrected_text="Beta",
                confidence=0.9,
            )
        ],
        [RecordingValidator("first_validator", seen)],
        seen,
    )
    second = RecordingModule(
        "second",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="second",
                module_key="second",
                id=1,
                original_text="Beta",
                corrected_text="Gamma",
                confidence=0.9,
            )
        ],
        [RecordingValidator("second_validator", seen)],
        seen,
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[
            ModuleRunSpec(instance_name="first", module_key="first", module=first),
            ModuleRunSpec(instance_name="second", module_key="second", module=second),
        ],
        config=AuditaConfig.from_sources(env={}),
        run_dir=tmp_path / "run",
    )

    assert seen[0] == ("propose", ["Alpha."])
    assert seen[1] == ("first_validator", ["Beta"])
    assert seen[2] == ("propose", ["Beta."])
    assert seen[3] == ("second_validator", ["Gamma"])
    assert result.transcript[0].text == "Gamma."
    assert len(result.applied_changes) == 2


def test_pipeline_runner_validator_order_respects_survivors(tmp_path):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Hello."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Hello"
            category: noun
            summary: "Hello."
        """
    )
    seen = []
    module = RecordingModule(
        "mod",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="mod",
                module_key="mod",
                id=1,
                original_text="Hello",
                corrected_text="Goodbye",
                confidence=0.9,
            )
        ],
        [
            RecordingValidator("first", seen, approve=False),
            RecordingLLMValidator("second", seen, approve=True),
        ],
        seen,
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="mod", module_key="mod", module=module)],
        config=AuditaConfig.from_sources(env={}),
        run_dir=tmp_path / "run",
    )

    assert ("first", ["Goodbye"]) in seen
    assert all(entry[0] != "second" for entry in seen)
    assert result.module_reports[0].validators[0].rejected_count == 1
    assert result.module_reports[0].validators[1].candidate_count == 0
    assert result.skipped_corrections[0].source == "validator:first"


def test_pipeline_runner_supports_deterministic_and_llm_validators_in_one_chain(tmp_path):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Alpha"
            category: noun
            summary: "Alpha."
        """
    )
    seen = []
    module = RecordingModule(
        "mixed",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="mixed",
                module_key="mixed",
                id=1,
                original_text="Alpha",
                corrected_text="Beta",
                confidence=0.9,
            )
        ],
        [
            RecordingValidator("deterministic_guard", seen),
            RecordingLLMValidator("llm_review", seen),
        ],
        seen,
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="mixed", module_key="mixed", module=module)],
        config=AuditaConfig.from_sources(env={}),
        run_dir=tmp_path / "run",
    )

    assert result.transcript[0].text == "Beta."
    assert [report.execution_kind for report in result.module_reports[0].validators] == [
        "deterministic",
        "llm",
    ]


def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the dam."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Jesters"
            category: faction
            summary: "The Jesters are a faction."
        """
    )
    module = RecordingModule(
        "mixed_real",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="mixed_real",
                module_key="glossary",
                id=1,
                original_text="gestures",
                corrected_text="Jesters",
                confidence=0.9,
            )
        ],
        [
            ProposalConfidenceValidator("proposal_confidence_guard", "glossary_confidence_threshold"),
            ProtectedGlossaryTermsValidator("protected_glossary_guard"),
            SpokenFormPlausibilityValidator("spoken_form_plausibility_review"),
            MeaningReversalValidator("meaning_reversal_review"),
        ],
        [],
    )
    llm_client = FakeStructuredLLMClient(
        {
            "mixed_real:spoken_form_plausibility_review": {
                "validations": [
                    {
                        "correction_index": 0,
                        "approved": True,
                        "confidence": 0.95,
                        "reason": "Likely phonetic mistranscription in context.",
                    }
                ]
            },
            "mixed_real:meaning_reversal_review": {
                "validations": [
                    {
                        "correction_index": 0,
                        "approved": True,
                        "confidence": 0.98,
                        "reason": "Does not reverse the segment meaning.",
                    }
                ]
            },
        }
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="mixed_real", module_key="glossary", module=module)],
        config=AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}),
        run_dir=tmp_path / "run",
        llm_client=llm_client,
    )

    assert result.transcript[0].text == "There were Jesters at the dam."
    assert [report.execution_kind for report in result.module_reports[0].validators] == [
        "deterministic",
        "deterministic",
        "llm",
        "llm",
    ]
    assert {call["stage_name"] for call in llm_client.calls} == {
        "mixed_real:spoken_form_plausibility_review",
        "mixed_real:meaning_reversal_review",
    }


def test_pipeline_runner_uses_real_protected_glossary_validator(tmp_path):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Hrank moves."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Hrank"
            category: pc
            summary: "Hrank is a player character."
        """
    )
    module = RecordingModule(
        "protected",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="protected",
                module_key="protected",
                id=1,
                original_text="Hrank",
                corrected_text="Frank",
                confidence=0.9,
            )
        ],
        [ProtectedGlossaryTermsValidator("protected_glossary_guard")],
        [],
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="protected", module_key="protected", module=module)],
        config=AuditaConfig.from_sources(env={}),
        run_dir=tmp_path / "run",
    )

    assert result.transcript[0].text == "Hrank moves."
    assert result.skipped_corrections[0].source == "validator:protected_glossary_guard"
    assert result.skipped_corrections[0].reason == "correction changes protected glossary term usage"


def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_section_order(tmp_path, monkeypatch):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."},
          {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "Beta."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Alpha"
            category: noun
            summary: "Alpha."
        """
    )
    sections = [
        TranscriptSection(
            section_index=0,
            start_index=0,
            segments=[IndexedSegment(index=0, segment=transcript[0])],
            token_count=1,
        ),
        TranscriptSection(
            section_index=1,
            start_index=1,
            segments=[IndexedSegment(index=1, segment=transcript[1])],
            token_count=1,
        ),
    ]
    seen = []
    module = ConcurrentRecordingModule(seen, threading.Barrier(2, timeout=1.0))
    monkeypatch.setattr(
        "audita.framework.runner.chunk_transcript",
        lambda working, max_tokens, min_section_tokens=1, target_section_count=None, exact_target_section_count=None: sections,
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="concurrent", module_key="concurrent", module=module)],
        config=AuditaConfig.from_sources(env={}, overrides=ConfigOverrides(llm_concurrency=2)),
        run_dir=tmp_path / "run",
    )

    assert [change.id for change in result.applied_changes] == [1, 2]
    assert result.applied_changes[0].corrected_text == "Alpha revised"
    assert result.applied_changes[1].corrected_text == "Beta revised"
    assert result.transcript[0].text == "Alpha revised."
    assert result.transcript[1].text == "Beta revised."


def test_pipeline_runner_passes_exact_target_sections_to_chunker(tmp_path, monkeypatch):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Alpha"
            category: noun
            summary: "Alpha."
        """
    )
    seen_args = {}
    sections = [
        TranscriptSection(
            section_index=0,
            start_index=0,
            segments=[IndexedSegment(index=0, segment=transcript[0])],
            token_count=1,
        )
    ]
    module = RecordingModule("noop", [], [], [])

    def _fake_chunk_transcript(working, max_tokens, min_section_tokens=1, target_section_count=None, exact_target_section_count=None):
        seen_args["target_section_count"] = target_section_count
        seen_args["exact_target_section_count"] = exact_target_section_count
        return sections

    monkeypatch.setattr("audita.framework.runner.chunk_transcript", _fake_chunk_transcript)

    runner = PipelineRunner()
    runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="noop", module_key="noop", module=module)],
        config=AuditaConfig.from_sources(env={}, overrides=ConfigOverrides(target_sections=3)),
        run_dir=tmp_path / "run",
    )

    assert seen_args == {
        "target_section_count": None,
        "exact_target_section_count": 3,
    }


def test_pipeline_runner_reports_first_llm_validator_rejection_in_chain_order(tmp_path):
    transcript = parse_transcript_json(
        """
        [
          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the dam."}
        ]
        """
    )
    glossary = parse_glossary_yaml(
        """
        glossary:
          - name: "Jesters"
            category: faction
            summary: "The Jesters are a faction."
        """
    )
    module = RecordingModule(
        "ordered_real",
        [
            CorrectionProposal(
                proposal_index=0,
                module_instance="ordered_real",
                module_key="glossary",
                id=1,
                original_text="gestures",
                corrected_text="Jesters",
                confidence=0.9,
            )
        ],
        [
            ProposalConfidenceValidator("proposal_confidence_guard", "glossary_confidence_threshold"),
            ProtectedGlossaryTermsValidator("protected_glossary_guard"),
            SpokenFormPlausibilityValidator("spoken_form_plausibility_review"),
            MeaningReversalValidator("meaning_reversal_review"),
        ],
        [],
    )
    llm_client = FakeStructuredLLMClient(
        {
            "ordered_real:spoken_form_plausibility_review": {
                "validations": [
                    {
                        "correction_index": 0,
                        "approved": False,
                        "confidence": 0.95,
                        "reason": "Not plausibly supported by spoken-form context.",
                    }
                ]
            },
            "ordered_real:meaning_reversal_review": {
                "validations": [
                    {
                        "correction_index": 0,
                        "approved": False,
                        "confidence": 0.98,
                        "reason": "Changes meaning too much.",
                    }
                ]
            },
        }
    )

    runner = PipelineRunner()
    result = runner.run(
        transcript=transcript,
        glossary=glossary,
        module_specs=[ModuleRunSpec(instance_name="ordered_real", module_key="glossary", module=module)],
        config=AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}),
        run_dir=tmp_path / "run",
        llm_client=llm_client,
    )

    assert result.skipped_corrections[0].source == "validator:spoken_form_plausibility_review"
    assert result.skipped_corrections[0].reason == "Not plausibly supported by spoken-form context."