"""
Prompt Template Base Classes

Provides templating infrastructure for RAG prompts with
support for variable substitution and validation.
"""

import re
from dataclasses import dataclass, field
from typing import Any


@dataclass
class PromptVariable:
    """Definition of a variable that can be used in a template."""

    name: str
    description: str
    required: bool = True
    default: str | None = None

    def validate(self, value: Any) -> tuple[bool, str | None]:
        """
        Validate a value for this variable.

        Returns:
            Tuple of (is_valid, error_message)
        """
        if self.required and value is None:
            return False, f"Required variable '{self.name}' is missing"

        if value is not None and not isinstance(value, str):
            return False, f"Variable '{self.name}' must be a string, got {type(value)}"

        return True, None


@dataclass
class PromptTemplate:
    """
    Template for generating LLM prompts with variable substitution.

    Supports simple {{variable}} syntax for substitution and
    validates that all required variables are provided.

    Example:
        template = PromptTemplate(
            name="qa_with_context",
            template=\"\"\"
            Context:
            {{context}}

            Question: {{question}}

            Answer based on the context provided.
            \"\"\",
            variables=[
                PromptVariable("context", "Retrieved document chunks"),
                PromptVariable("question", "User's question"),
            ]
        )

        prompt = template.render(
            context="The critical mass of U-235 is approximately 52 kg...",
            question="What is the critical mass of U-235?"
        )
    """

    name: str
    template: str
    description: str = ""
    variables: list[PromptVariable] = field(default_factory=list)

    # Template metadata
    category: str = "general"
    version: str = "1.0"

    def __post_init__(self):
        """Extract variables from template if not provided."""
        if not self.variables:
            self.variables = self._extract_variables()

    def _extract_variables(self) -> list[PromptVariable]:
        """Extract variable names from template."""
        pattern = r"\{\{(\w+)\}\}"
        matches = re.findall(pattern, self.template)
        return [PromptVariable(name=name, description="") for name in set(matches)]

    def validate_inputs(self, **kwargs) -> tuple[bool, list[str]]:
        """
        Validate inputs against template variables.

        Returns:
            Tuple of (is_valid, list_of_errors)
        """
        errors = []
        provided_vars = set(kwargs.keys())
        required_vars = {v.name for v in self.variables if v.required}

        # Check for missing required variables
        missing = required_vars - provided_vars
        if missing:
            errors.extend([f"Missing required variable: {v}" for v in missing])

        # Validate each variable
        for var in self.variables:
            value = kwargs.get(var.name, var.default)
            is_valid, error = var.validate(value)
            if not is_valid and error:
                errors.append(error)

        return len(errors) == 0, errors

    def render(self, **kwargs) -> str:
        """
        Render the template with provided variables.

        Args:
            **kwargs: Variable values to substitute

        Returns:
            Rendered prompt string

        Raises:
            ValueError: If required variables are missing or invalid
        """
        is_valid, errors = self.validate_inputs(**kwargs)
        if not is_valid:
            raise ValueError(f"Invalid template inputs: {'; '.join(errors)}")

        # Apply defaults for optional variables
        for var in self.variables:
            if var.name not in kwargs and var.default is not None:
                kwargs[var.name] = var.default

        # Perform substitution
        result = self.template
        for name, value in kwargs.items():
            if value is not None:
                result = result.replace(f"{{{{{name}}}}}", str(value))

        return result.strip()

    def get_variable_names(self) -> list[str]:
        """Get list of variable names in this template."""
        return [v.name for v in self.variables]


@dataclass
class SystemPrompt:
    """
    System prompt configuration for a specific task type.

    System prompts define the AI's role and constraints for
    a particular use case (e.g., NCS Q&A, document summarization).
    """

    name: str
    prompt: str
    description: str = ""

    # Task constraints
    max_input_tokens: int = 8000
    max_output_tokens: int = 4096
    recommended_temperature: float = 0.3

    # Safety settings
    require_citations: bool = True
    allow_uncertainty: bool = True  # Can say "I don't know"
    conservative_mode: bool = True  # Prefer conservative answers for safety

    def to_string(self) -> str:
        """Get the prompt as a string."""
        return self.prompt.strip()


class PromptRegistry:
    """
    Registry for managing prompt templates.

    Provides centralized access to all templates with
    versioning and validation support.
    """

    def __init__(self):
        """Initialize the registry."""
        self._templates: dict[str, PromptTemplate] = {}
        self._system_prompts: dict[str, SystemPrompt] = {}

    def register_template(self, template: PromptTemplate) -> None:
        """Register a prompt template."""
        key = f"{template.category}:{template.name}"
        self._templates[key] = template

    def register_system_prompt(self, system_prompt: SystemPrompt) -> None:
        """Register a system prompt."""
        self._system_prompts[system_prompt.name] = system_prompt

    def get_template(
        self, name: str, category: str = "general"
    ) -> PromptTemplate | None:
        """Get a template by name and category."""
        key = f"{category}:{name}"
        return self._templates.get(key)

    def get_system_prompt(self, name: str) -> SystemPrompt | None:
        """Get a system prompt by name."""
        return self._system_prompts.get(name)

    def list_templates(self, category: str | None = None) -> list[PromptTemplate]:
        """List all templates, optionally filtered by category."""
        templates = list(self._templates.values())
        if category:
            templates = [t for t in templates if t.category == category]
        return templates

    def list_system_prompts(self) -> list[SystemPrompt]:
        """List all system prompts."""
        return list(self._system_prompts.values())


# Global registry instance
_registry = PromptRegistry()


def get_registry() -> PromptRegistry:
    """Get the global prompt registry."""
    return _registry


def register_template(template: PromptTemplate) -> PromptTemplate:
    """Register a template with the global registry."""
    _registry.register_template(template)
    return template


def register_system_prompt(system_prompt: SystemPrompt) -> SystemPrompt:
    """Register a system prompt with the global registry."""
    _registry.register_system_prompt(system_prompt)
    return system_prompt
