"""Configuration for LLM processor in MMIRAGE."""

from dataclasses import dataclass, field

import logging
from typing import Dict, Optional, Sequence, Type, Any, List
from pydantic import BaseModel, create_model

from mmirage.core.process.variables import BaseVar, OutputVar
from sglang.srt.server_args import ServerArgs

from mmirage.core.process.base import BaseProcessorConfig
from jinja2 import Environment, meta

logger = logging.getLogger(__name__)
env = Environment()


@dataclass
class SGLangLLMConfig(BaseProcessorConfig):
    """Configuration for LLM processor using SGLang.

    Attributes:
        type: Type identifier (must be "llm").
        server_args: SGLang server arguments including model path and TP size.
        default_sampling_params: Default sampling parameters for generation.
    """

    server_args: ServerArgs = field(
        default_factory=lambda: ServerArgs(model_path="none")
    )
    default_sampling_params: Dict[str, Any] = field(default_factory=dict)


@dataclass
class LLMOutputVar(OutputVar):
    """Output variable generated by LLM processor.

    Uses Jinja2 templating for prompts and supports both plain text
    and structured JSON outputs.

    Attributes:
        name: Name of the variable.
        type: Type identifier (must be "llm").
        prompt: Jinja2 template for the LLM prompt.
        output_schema: List of field names for JSON output (empty for plain text).
        output_type: Output format - "JSON" or "plain".
    """

    prompt: str = ""
    output_schema: List[str] = field(default_factory=list)
    output_type: str = ""

    def get_output_schema(self) -> Optional[Type[BaseModel]]:
        """Generate a Pydantic model for JSON output validation.

        Returns:
            A Pydantic BaseModel class if output_type is "JSON" and
            output_schema is non-empty, otherwise None.
        """
        if self.output_type == "JSON" and self.output_schema:
            fields: Dict[str, Any] = {var: (str, ...) for var in self.output_schema}
            return create_model("OutputSchema", **fields)
        return None

    def is_computable(self, vars: Sequence[BaseVar]) -> bool:
        """Check if all variables referenced in the prompt are available.

        Args:
            vars: Sequence of currently available variables.

        Returns:
            True if all template variables are declared, False otherwise.
        """
        parsed_content = env.parse(self.prompt)
        template_vars = meta.find_undeclared_variables(parsed_content)

        var_names = set(map(lambda v: v.name, vars))
        undeclared_vars = template_vars - var_names

        if len(undeclared_vars) > 0:
            logger.info(
                f"⚠️ Undeclared variables found for {self.name}: {undeclared_vars}"
            )
            return False

        return True
