import json
from pathlib import Path
from typing import List, Protocol

from .chunking import TranscriptSection
from .config import AuditaConfig
from .prompts import build_glossary_correction_messages, build_grammar_correction_messages
from .schemas import CorrectionCandidate, CorrectionSet, Glossary


class LLMClient(Protocol):
    def create_corrections(self, messages: List[dict], config: AuditaConfig) -> CorrectionSet:
        ...


class CorrectionPass(Protocol):
    def run(
        self,
        section: TranscriptSection,
        glossary: Glossary,
        config: AuditaConfig,
        run_dir: Path,
        retry_pass: bool = False,
    ) -> List[CorrectionCandidate]:
        ...


class GlossaryCorrectionPass:
    def __init__(self, llm_client: LLMClient) -> None:
        self._llm_client = llm_client

    def run(
        self,
        section: TranscriptSection,
        glossary: Glossary,
        config: AuditaConfig,
        run_dir: Path,
        retry_pass: bool = False,
    ) -> List[CorrectionCandidate]:
        messages = build_glossary_correction_messages(section, glossary, retry_pass=retry_pass)
        prompt_path = run_dir / f"prompt-{section.section_index:04d}.json"
        prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

        response = self._llm_client.create_corrections(messages, config)
        response_path = run_dir / f"corrections-{section.section_index:04d}.json"
        response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8")
        return list(response.corrections)


class GrammarCorrectionPass:
    def __init__(self, llm_client: LLMClient) -> None:
        self._llm_client = llm_client

    def run(
        self,
        section: TranscriptSection,
        glossary: Glossary,
        config: AuditaConfig,
        run_dir: Path,
        retry_pass: bool = False,
    ) -> List[CorrectionCandidate]:
        messages = build_grammar_correction_messages(section, glossary, retry_pass=retry_pass)
        prompt_path = run_dir / f"prompt-{section.section_index:04d}.json"
        prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

        response = self._llm_client.create_corrections(messages, config)
        response_path = run_dir / f"corrections-{section.section_index:04d}.json"
        response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8")
        return list(response.corrections)