"""DSPy module for optimizing Pydantic field descriptions and prompts."""

import re
from typing import Any

import dspy


class OptimizeFieldDescription(dspy.Signature):
    """Rewrite the given field description to be clearer, more specific, \
and more effective for guiding structured data extraction. \
Output ONLY the improved description text."""

    field_description: str = dspy.InputField(
        desc="The current field description to improve"
    )
    field_type: str = dspy.InputField(desc="The data type of the field")
    optimized_field_description: str = dspy.OutputField(
        desc="The improved field description"
    )


class OptimizeSystemPrompt(dspy.Signature):
    """Rewrite the given system prompt to be clearer and more effective \
for guiding structured data extraction. \
Output ONLY the improved system prompt text."""

    system_prompt: str = dspy.InputField(
        desc="The current system prompt to improve"
    )
    optimized_system_prompt: str = dspy.OutputField(
        desc="The improved system prompt"
    )


class OptimizeInstructionPrompt(dspy.Signature):
    """Rewrite the given instruction prompt to be clearer and more \
effective for guiding structured data extraction. \
Preserve any template placeholders exactly as they appear. \
Output ONLY the improved instruction prompt text."""

    instruction_prompt: str = dspy.InputField(
        desc="The current instruction prompt to improve"
    )
    optimized_instruction_prompt: str = dspy.OutputField(
        desc="The improved instruction prompt"
    )


def _make_contextual_field_signature(
    model_name: str,
) -> type[dspy.Signature]:
    """Create a context-aware signature for field description optimization.

    Uses three lightweight techniques to give MiPROV2's proposer domain context
    without bloating token usage:
    1. Dynamic class name (0 extra tokens — replaces generic name)
    2. Concise docstring with model name (~25 tokens)
    3. field_name input field (~2-5 tokens per call)

    Args:
        model_name: Name of the Pydantic model being optimized.

    Returns:
        A Signature class with contextual class name, docstring, and fields.
    """
    docstring = (
        f"Improve a field description for {model_name} structured data extraction. "
        f"Output ONLY the improved description — a short descriptive phrase, "
        f"not instructions."
    )

    # Dynamic class name gives the proposer domain signal for free
    cls = type(
        f"Optimize{model_name}FieldDescription",
        (dspy.Signature,),
        {
            "__doc__": docstring,
            "__annotations__": {
                "field_name": str,
                "field_description": str,
                "field_type": str,
                "optimized_field_description": str,
            },
            "field_name": dspy.InputField(desc="Name of the field being optimized"),
            "field_description": dspy.InputField(
                desc="The current field description to improve"
            ),
            "field_type": dspy.InputField(desc="The data type of the field"),
            "optimized_field_description": dspy.OutputField(
                desc="The improved field description — a short descriptive phrase"
            ),
        },
    )
    return cls


def _make_contextual_system_prompt_signature(
    model_name: str,
) -> type[dspy.Signature]:
    """Create a context-aware signature for system prompt optimization."""
    docstring = (
        f"Improve the system prompt for {model_name} structured data extraction. "
        f"Output ONLY the improved system prompt text."
    )

    cls = type(
        f"Optimize{model_name}SystemPrompt",
        (dspy.Signature,),
        {
            "__doc__": docstring,
            "__annotations__": {
                "system_prompt": str,
                "optimized_system_prompt": str,
            },
            "system_prompt": dspy.InputField(
                desc="The current system prompt to improve"
            ),
            "optimized_system_prompt": dspy.OutputField(
                desc="The improved system prompt"
            ),
        },
    )
    return cls


def _make_contextual_instruction_prompt_signature(
    model_name: str,
) -> type[dspy.Signature]:
    """Create a context-aware signature for instruction prompt optimization."""
    docstring = (
        f"Improve the instruction prompt for {model_name} structured data extraction. "
        f"Preserve any template placeholders exactly as they appear. "
        f"Output ONLY the improved instruction prompt text."
    )

    cls = type(
        f"Optimize{model_name}InstructionPrompt",
        (dspy.Signature,),
        {
            "__doc__": docstring,
            "__annotations__": {
                "instruction_prompt": str,
                "optimized_instruction_prompt": str,
            },
            "instruction_prompt": dspy.InputField(
                desc="The current instruction prompt to improve"
            ),
            "optimized_instruction_prompt": dspy.OutputField(
                desc="The improved instruction prompt"
            ),
        },
    )
    return cls


class PydanticOptimizerModule(dspy.Module):
    """DSPy module for optimizing field descriptions, system prompts, and instruction prompts."""

    def __init__(
        self,
        field_descriptions: dict[str, str] | None = None,
        field_types: dict[str, str] | None = None,
        has_system_prompt: bool = False,
        has_instruction_prompt: bool = False,
        model_name: str | None = None,
    ):
        """Initialize the optimizer module.

        Args:
            field_descriptions: Dictionary mapping field paths to their descriptions.
            field_types: Dictionary mapping field paths to their type names as strings.
            has_system_prompt: Whether to optimize a system prompt.
            has_instruction_prompt: Whether to optimize an instruction prompt.
            model_name: Name of the Pydantic model (for contextual signatures).
        """
        super().__init__()

        # Store field descriptions and types for optimization
        self.field_descriptions = field_descriptions or {}
        self.field_types = field_types or {}

        # Choose signature: contextual (with model name) or generic
        if model_name:
            field_sig = _make_contextual_field_signature(model_name)
            self._has_field_name_input = True
        else:
            field_sig = OptimizeFieldDescription
            self._has_field_name_input = False

        # Create optimizers for each field description
        self.field_optimizers: dict[str, dspy.ChainOfThought] = {}
        for field_path, description in self.field_descriptions.items():
            self.field_optimizers[field_path] = dspy.ChainOfThought(field_sig)

        # Create optimizers for prompts if needed
        self.has_system_prompt = has_system_prompt
        self.has_instruction_prompt = has_instruction_prompt

        if has_system_prompt:
            if model_name:
                sys_sig = _make_contextual_system_prompt_signature(model_name)
            else:
                sys_sig = OptimizeSystemPrompt
            self.system_prompt_optimizer = dspy.ChainOfThought(sys_sig)

        if has_instruction_prompt:
            if model_name:
                instr_sig = _make_contextual_instruction_prompt_signature(model_name)
            else:
                instr_sig = OptimizeInstructionPrompt
            self.instruction_prompt_optimizer = dspy.ChainOfThought(instr_sig)

    def _remove_optimization_markers(self, text: str) -> str:
        """Remove optimization instruction markers from optimized prompt text."""
        text = re.sub(
            r"ORIGINAL_PROMPT_TEMPLATE:.*?PROMPT_TEMPLATE_REWRITE_INSTRUCTIONS:",
            "",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        )
        text = re.sub(
            r"PROMPT_TEMPLATE_REWRITE_INSTRUCTIONS:.*$",
            "",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        )
        text = re.sub(
            r"\n\n(?:ORIGINAL_PROMPT_TEMPLATE|PROMPT_TEMPLATE_REWRITE_INSTRUCTIONS):.*?\.?",
            "",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        )
        return text.strip()

    def forward(
        self,
        system_prompt: str | None = None,
        instruction_prompt: str | None = None,
        **field_descriptions: str,
    ) -> dict[str, Any]:
        """Forward pass for optimization.

        Args:
            system_prompt: System prompt to optimize (if provided).
            instruction_prompt: Instruction prompt to optimize (if provided).
            **field_descriptions: Field descriptions to optimize (keyed by field path).

        Returns:
            Dictionary with optimized field descriptions and prompts.
        """
        optimized: dict[str, Any] = {}

        # Optimize field descriptions
        for field_path, description in field_descriptions.items():
            # Skip field_type_ prefixed entries (they're metadata, not descriptions)
            if field_path.startswith("field_type_"):
                continue

            if field_path in self.field_optimizers:
                optimizer = self.field_optimizers[field_path]
                # Get field type from keyword arguments (passed as field_type_{field_path})
                # or fall back to stored field_types
                field_type_key = f"field_type_{field_path}"
                field_type = field_descriptions.get(field_type_key, "")
                if not field_type:
                    field_type = self.field_types.get(field_path, "")
                # Pass field_name if the signature accepts it (contextual signatures do)
                kwargs: dict[str, str] = {
                    "field_description": description,
                    "field_type": field_type,
                }
                if self._has_field_name_input:
                    kwargs["field_name"] = field_path
                result = optimizer(**kwargs)
                optimized[f"optimized_{field_path}"] = (
                    result.optimized_field_description
                )

        # Optimize system prompt
        if self.has_system_prompt and system_prompt is not None:
            result = self.system_prompt_optimizer(system_prompt=system_prompt)
            optimized["optimized_system_prompt"] = result.optimized_system_prompt

        # Optimize instruction prompt
        if self.has_instruction_prompt and instruction_prompt is not None:
            # Check if instruction prompt contains template placeholders
            placeholder_pattern = r"\{([^}]+)\}"
            placeholders = re.findall(placeholder_pattern, instruction_prompt)

            if placeholders:
                # Extract unique placeholders
                unique_placeholders = list(dict.fromkeys(placeholders))
                placeholder_list = ", ".join([f"{{{p}}}" for p in unique_placeholders])

                # Create optimization prompt with clear instructions
                original_prompt_marker = "ORIGINAL_PROMPT_TEMPLATE"
                rewrite_instructions_marker = "PROMPT_TEMPLATE_REWRITE_INSTRUCTIONS"

                # Build field descriptions and types if available
                field_descriptions_text = ""
                if field_descriptions:
                    field_desc_lines = []
                    for key, value in field_descriptions.items():
                        # Skip field_type_ prefixed entries (they're metadata, not descriptions)
                        if key.startswith("field_type_"):
                            continue
                        # Get field type if available (from kwargs or stored types)
                        field_type_key = f"field_type_{key}"
                        field_type = field_descriptions.get(field_type_key, "")
                        if not field_type:
                            field_type = self.field_types.get(key, "")
                        if field_type:
                            field_desc_lines.append(f"- {key} ({field_type}): {value}")
                        else:
                            field_desc_lines.append(f"- {key}: {value}")
                    if field_desc_lines:
                        field_descriptions_text = (
                            "\n\nAdditionally, here are the field descriptions and types "
                            "for the structured information you are helping to extract:"
                            "\n" + "\n".join(field_desc_lines)
                        )

                optimization_prompt = (
                    f"{original_prompt_marker}:\n{instruction_prompt}\n\n"
                    f"{rewrite_instructions_marker}:\n"
                    f"You are provided with a {original_prompt_marker} which is used for structured data "
                    f"extraction.\n\n"
                    f"This template contains placeholders ({placeholder_list}) that will be "
                    f"replaced at runtime. {field_descriptions_text}\n\n"
                    f"You may optimize the template for clarity, logical flow, and effectiveness, "
                    f"including reordering, restructuring, or rephrasing text, while preserving the placeholders.\n\n"
                )

                # Optimize the prompt directly with placeholders intact
                result = self.instruction_prompt_optimizer(instruction_prompt=optimization_prompt)
                optimized_prompt = result.optimized_instruction_prompt

                # Remove the instruction markers and their content
                optimized_prompt = self._remove_optimization_markers(optimized_prompt)

                # Verify all placeholders are present exactly once
                for placeholder_name in unique_placeholders:
                    placeholder_str = f"{{{placeholder_name}}}"
                    count = optimized_prompt.count(placeholder_str)

                    if count == 0:
                        break
                    elif count > 1:
                        # Duplicate - keep only the first occurrence
                        first_idx = optimized_prompt.find(placeholder_str)
                        if first_idx != -1:
                            before_first = optimized_prompt[: first_idx + len(placeholder_str)]
                            after_first = optimized_prompt[first_idx + len(placeholder_str) :]
                            after_first = after_first.replace(placeholder_str, "")
                            optimized_prompt = before_first + after_first

                # Final check: if any placeholders are still missing, use original prompt
                all_placeholders_present = all(
                    f"{{{p}}}" in optimized_prompt for p in unique_placeholders
                )
                if not all_placeholders_present:
                    optimized_prompt = instruction_prompt

                optimized["optimized_instruction_prompt"] = optimized_prompt.strip()
            else:
                result = self.instruction_prompt_optimizer(instruction_prompt=instruction_prompt)
                optimized["optimized_instruction_prompt"] = result.optimized_instruction_prompt

        return dspy.Prediction(**optimized)
