import importlib
import inspect
import pkgutil
import time
from typing import ClassVar

from pydantic import BaseModel, Field

from yaduha.agent import Agent
from yaduha.language.language import Language
from yaduha.loader import LanguageLoader
from yaduha.translator import BackTranslation
from yaduha.translator.back_translator import BackTranslator


class BackTranslationResponse(BaseModel):
    """Structured LLM response for back-translation with grammaticality judgment.

    Fields are ordered so that `reasoning` comes first — since structured output
    generates fields sequentially, this forces chain-of-thought before the answer.
    """

    reasoning: str = Field(
        ...,
        description=(
            "Step-by-step reasoning explaining how you analyzed the sentence "
            "by tracing through the source code. Reference specific functions, "
            "dictionaries, and suffix rules from the code."
        ),
    )
    grammatical: bool = Field(
        ...,
        description=(
            "Whether this sentence could have been produced by the sentence "
            "models in the source code. True means every part of the sentence "
            "is consistent with the code's output patterns."
        ),
    )
    english_translation: str | None = Field(
        None,
        description=(
            "A natural English translation of the sentence. "
            "Provide this when the sentence is grammatical or partially understandable. "
            "Set to null only when the sentence is completely unintelligible."
        ),
    )


_PREAMBLE = """\
You are analyzing sentences in a constructed/low-resource language by reading \
the Python source code that generates them.

Below is the COMPLETE source code of the language package. The code defines:
- Vocabulary (nouns, verbs, etc.) as VocabEntry(english, target) pairs
- Grammatical enumerations (e.g. proximity, plurality, tense/aspect, pronouns)
- Morphology rules (suffix functions, consonant transformations)
- Sentence models (Pydantic classes) whose __str__ methods render the surface form

Your task: Given a sentence in the target language, reverse-engineer it by \
tracing through the code.
1. Identify which sentence model could produce it based on word count and \
morphological patterns.
2. Look up vocabulary items by matching target forms from the vocab module.
3. Identify suffixes by checking the suffix functions in the code.
4. For transitive sentences, check for consonant transformations on the verb \
stem and any pronoun prefixes.
5. Determine if ALL parts of the sentence are consistent with the code. \
If any part cannot be explained by the code, mark as ungrammatical.
6. If grammatical, compose a natural English translation from the identified \
components.

IMPORTANT: The surface form uses hyphens (-) to separate morphemes.
"""


def _get_package_sources(language: Language) -> dict[str, str]:
    """Extract source code from all modules in a language package."""
    # Find the module where the first sentence type is defined
    sentence_type = language.sentence_types[0]
    module = inspect.getmodule(sentence_type)
    if module is None:
        raise ValueError(f"Cannot find module for sentence type {sentence_type.__name__}")

    package_name = module.__package__
    sources: dict[str, str] = {}

    # Get the main module source
    sources[module.__name__] = inspect.getsource(module)

    # Walk submodules to get vocab, prompts, etc.
    if package_name:
        package_module = importlib.import_module(package_name)
        if hasattr(package_module, "__path__"):
            for _importer, modname, _ispkg in pkgutil.walk_packages(
                package_module.__path__, prefix=package_name + "."
            ):
                try:
                    submod = importlib.import_module(modname)
                    sources[modname] = inspect.getsource(submod)
                except Exception:
                    continue

    return sources


def _build_examples_section(language: Language, n_examples: int) -> str:
    """Generate examples showing (surface_form, json_structure) pairs."""
    lines = [
        "## Examples (generated by running the code above)\n",
        "Each example shows the surface form produced by __str__ and the "
        "JSON structure from model_dump_json().\n",
    ]

    example_num = 0
    for sentence_type in language.sentence_types:
        type_name = sentence_type.__name__

        # Use sample_iter if available (language-specific), otherwise get_examples
        if hasattr(sentence_type, "sample_iter"):
            per_type = max(1, n_examples // len(language.sentence_types))
            sentences = list(sentence_type.sample_iter(per_type))  # type: ignore[attr-defined]
        else:
            sentences = [ex[1] for ex in sentence_type.get_examples()]

        for sentence in sentences:
            example_num += 1
            lines.append(f"### Example {example_num} ({type_name})")
            lines.append(f"- Surface form: `{sentence}`")
            lines.append(f"- JSON: `{sentence.model_dump_json()}`")
            lines.append("")

    return "\n".join(lines)


def _build_source_code_instructions(language: Language, n_examples: int = 10) -> str:
    """Build the full system prompt from language source code and examples."""
    sources = _get_package_sources(language)

    prompt = _PREAMBLE + "\n"

    for module_name, source in sources.items():
        prompt += f"## Source Code: {module_name}\n"
        prompt += f"```python\n{source}\n```\n\n"

    prompt += _build_examples_section(language, n_examples)

    return prompt


class InstructionsBackTranslator(BackTranslator):
    """Back-translate target-language text using source code as instructions.

    The LLM receives the actual Python source code of the language package
    and uses it to reverse-engineer sentences, determining both grammaticality
    and meaning.
    """

    name: ClassVar[str] = "instructions_back_translator"
    description: ClassVar[str] = (
        "Translate text from the target language back to the source language "
        "using the language package source code as instructions."
    )

    agent: Agent
    instructions: str

    @classmethod
    def from_language(
        cls,
        language_code: str,
        agent: Agent,
        n_examples: int = 10,
        **kwargs,
    ) -> "InstructionsBackTranslator":
        """Create an InstructionsBackTranslator from an installed language package.

        Loads the language, extracts source code via inspect.getsource(),
        and generates examples from the sentence types.

        Args:
            language_code: Language code (e.g., 'ovp')
            agent: Agent to use for back-translation
            n_examples: Number of generated examples to include in the prompt

        Returns:
            InstructionsBackTranslator instance
        """
        language = LanguageLoader.load_language(language_code)
        instructions = _build_source_code_instructions(language, n_examples)
        return cls(agent=agent, instructions=instructions, **kwargs)

    def back_translate(self, text: str) -> BackTranslation:
        start_time = time.time()
        response = self.agent.get_response(
            messages=[
                {"role": "system", "content": self.instructions},
                {"role": "user", "content": text},
            ],
            response_format=BackTranslationResponse,
        )
        translation_time = time.time() - start_time

        result: BackTranslationResponse = response.content

        self.log(
            data={
                "event": "back_translation_complete",
                "back_translator": self.name,
                "target": text,
                "source": result.english_translation or "",
                "grammatical": result.grammatical,
                "reasoning": result.reasoning,
                "translation_time": translation_time,
                "prompt_tokens": response.prompt_tokens,
                "completion_tokens": response.completion_tokens,
            }
        )

        return BackTranslation(
            source=result.english_translation or "",
            target=text,
            translation_time=translation_time,
            prompt_tokens=response.prompt_tokens,
            completion_tokens=response.completion_tokens,
        )
