# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import asyncio
import logging
from typing import Any, Dict, Generic, List, Optional, TypeVar

import openai
from mantisdk.algorithm.gepa.lib.api import optimize

from mantisdk.adapter import TraceAdapter
from mantisdk.adapter.messages import TraceToMessages
from mantisdk.algorithm.base import Algorithm
from mantisdk.algorithm.gepa.adapter import (
    MantisdkDataInst,
    MantisdkGEPAAdapter,
)
from mantisdk.algorithm.gepa.tracing import GEPATracingContext
from mantisdk.algorithm.utils import with_llm_proxy, with_store
from mantisdk.types import Dataset, PromptTemplate, TracingConfig

logger = logging.getLogger(__name__)

T_task = TypeVar("T_task")


TEMPLATE_AWARE_REFLECTION_PROMPT = """You are an expert at improving LLM prompts based on observed failures.

## Current Prompt Template
```
<curr_instructions>
```

## Observed Failures
The following examples show where the current prompt gave INCORRECT outputs. Study the pattern of failures carefully:
```
<inputs_outputs_feedback>
```

## Your Task
Write an IMPROVED prompt template that fixes the observed failures.

**CRITICAL REQUIREMENTS**:
1. **Analyze the failure pattern**: Look at WHAT the prompt got wrong. Is it too strict? Too lenient? Missing context? Misunderstanding the task?
2. **Make a CONCEPTUAL FIX**: Don't just tweak wording - fundamentally change the approach if the current strategy is flawed
3. **Preserve placeholders**: Keep any {variable_name} placeholders exactly as they appear (e.g., {session}, {input}, {output}, or {{session}}, {{input}}, {{output}})
4. **Be specific**: Add concrete criteria, examples, or decision rules based on what the failures reveal
5. Read the inputs carefully and identify the input format and infer detailed task description about the task I wish to solve with the assistant
6. Read all the assistant responses and the corresponding feedback. Identify all niche and domain specific factual information about the task and include it in the instruction, as a lot of it may not be available to the assistant in the future. The assistant may have utilized a generalizable strategy to solve the task, if so, include that in the instruction as well.

**Think step by step**:
- What is the PATTERN in the failures?
- WHY is the current prompt failing on these cases?
- What SPECIFIC change would fix this pattern?

Output your improved template within ``` blocks."""


class GEPA(Algorithm, Generic[T_task]):
    """GEPA (Genetic-Pareto) algorithm for Mantisdk.

    This algorithm optimizes prompt templates (and potentially other text resources)
    using an evolutionary approach with LLM-based reflection.
    
    GEPA maintains a population of candidate prompts and evolves them by:
    1. Evaluating candidates on training data
    2. Using an LLM to reflect on failures and propose improvements
    3. Selecting the best candidates based on validation performance
    """

    # Algorithm-owned tracing configuration.
    # Defines environment and tags for all traces generated by GEPA.
    # Note: session_id will be overridden per-run in the run() method
    TRACING_CONFIG = TracingConfig(
        environment="mantisdk-gepa",
        algorithm_name="gepa",
    )
    
    def _get_tracing_config_with_session(self, session_id: str) -> TracingConfig:
        """Create a TracingConfig with a specific session_id for this run."""
        return TracingConfig(
            environment="mantisdk-gepa",
            algorithm_name="gepa",
            session_id=session_id,
        )

    def __init__(
        self,
        *,
        max_metric_calls: int = 100,
        reflection_minibatch_size: int = 5,
        population_size: int = 4,
        adapter: Optional[TraceAdapter] = None,
        rollout_batch_timeout: float = 600.0,
        reflection_prompt_template: Optional[str] = None,
        reflection_metadata: Optional[Dict[str, Any]] = None,
        target_model_config: Optional[Dict[str, Any]] = None,
        **gepa_kwargs: Any,
    ) -> None:
        """Initialize the GEPA algorithm.

        Args:
            max_metric_calls: Maximum number of evaluations (budget).
            reflection_minibatch_size: Batch size for reflection.
            population_size: Size of the population in evolutionary search.
            adapter: TraceAdapter to convert spans to messages. Defaults to TraceToMessages.
            rollout_batch_timeout: Timeout for waiting for rollouts.
            reflection_prompt_template: Custom prompt template for reflection.
            reflection_metadata: Metadata for reflection LLM traces (Langfuse).
            target_model_config: Direct model configuration (apiKey, baseUrl) to bypass proxy if needed.
            **gepa_kwargs: Additional arguments passed to gepa.optimize.
        """
        super().__init__()
        self.max_metric_calls = max_metric_calls
        self.reflection_minibatch_size = reflection_minibatch_size
        self.population_size = population_size
        self.adapter = adapter or TraceToMessages()
        self.rollout_batch_timeout = rollout_batch_timeout
        self.reflection_metadata = reflection_metadata
        self.target_model_config = target_model_config
        
        # Set default reflection prompt template if not provided
        if reflection_prompt_template is not None:
            gepa_kwargs["reflection_prompt_template"] = reflection_prompt_template
        elif "reflection_prompt_template" not in gepa_kwargs:
            gepa_kwargs["reflection_prompt_template"] = TEMPLATE_AWARE_REFLECTION_PROMPT
            
        self.gepa_kwargs = gepa_kwargs
        self._best_candidate: Optional[Dict[str, str]] = None
        self._best_score: float = 0.0
        self._full_result: Optional[Any] = None  # Store full GEPAResult for history access

    def get_best_prompt(self) -> Optional[PromptTemplate]:
        """Get the best prompt found during optimization."""
        if self._best_candidate:
            # Return the first PromptTemplate value
            for key, value in self._best_candidate.items():
                return PromptTemplate(template=value, engine="f-string")
        return None

    @with_store
    @with_llm_proxy(required=True)
    async def run(
        self,
        llm_proxy,  # injected by decorator
        store,  # injected by decorator
        train_dataset: Optional[Dataset[T_task]] = None,
        val_dataset: Optional[Dataset[T_task]] = None,
    ) -> None:
        """Run the GEPA optimization loop.

        Args:
            train_dataset: Dataset used for optimization (training).
            val_dataset: Dataset used for validation.
        """
        if train_dataset is None:
            raise ValueError("train_dataset is required for GEPA optimization.")

        store = self.get_store()
        assert store is not None
        assert llm_proxy is not None

        loop = asyncio.get_running_loop()
        
        # Get initial resources to find the target resource to optimize
        initial_resources = self.get_initial_resources()
        if not initial_resources:
            raise ValueError("Initial resources must be set before running GEPA.")
        
        # Find the target resource to optimize
        target_resource_name = None
        initial_candidate: Dict[str, str] = {}
        
        for name, res in initial_resources.items():
            if isinstance(res, PromptTemplate):
                target_resource_name = name
                initial_candidate[name] = res.template
        
        if not target_resource_name:
            raise ValueError("No PromptTemplate found in initial resources to optimize.")

        logger.info(f"GEPA will optimize resource: {target_resource_name}")
        logger.info(f"Initial prompt: {initial_candidate[target_resource_name][:100]}...")

        # Setup Reflection LLM via Proxy
        llm_resource = llm_proxy.as_resource()
        
        # Create tracing context for detailed execution tracking
        # This generates a unique session_id for grouping all traces in this run
        tracing_context = GEPATracingContext()
        
        # Create a TracingConfig with this run's session_id
        run_tracing_config = self._get_tracing_config_with_session(tracing_context.session_id)
        
        logger.info(f"GEPA session started: {tracing_context.session_id}")
        
        # Create the bridge adapter with tracing config and context
        gepa_adapter = MantisdkGEPAAdapter(
            store=store,
            loop=loop,
            resource_name=target_resource_name,
            adapter=self.adapter,
            llm_proxy_resource=llm_resource,
            rollout_batch_timeout=self.rollout_batch_timeout,
            tracing_config=run_tracing_config,
            tracing_context=tracing_context,
        )

        # Import here (inside run) to avoid circular imports:
        # `mantisdk.algorithm.gepa.__init__` re-exports `GEPA` from this module.
        from mantisdk.algorithm import gepa as gepa_module
        
        @gepa_module.reflection
        def reflection_lm(prompt: str) -> str:
            """GEPA reflection LM callable - takes a string prompt, returns string response.
            
            Routes through the LLMProxy for consistent tracing and OpenTelemetry export.
            The proxy handles API keys and model routing - we just need to call it.
            
            IMPORTANT: The LLMProxy must be configured with callbacks: ["opentelemetry"]
            (not "return_token_ids") when calling OpenAI providers.
            
            The reflection phase is tagged distinctly from validation-eval to enable
            proper filtering in the Mantis UI. Tags include:
            - "gepa" (algorithm)
            - "reflection" (phase)
            - "llm-reflection" (explicit marker)
            - "gen-N" (generation number)
            """
            import litellm
            
            # Track phase transition: entering reflection, then starting new generation
            tracing_context.set_phase("reflection")
            
            logger.info(f"[reflection_lm] === REFLECTION CALLED (gen-{tracing_context.generation}) ===")
            logger.info(f"[reflection_lm] Session: {tracing_context.session_id}")
            logger.info(f"[reflection_lm] Model: {llm_resource.model}")
            logger.info(f"[reflection_lm] Proxy Endpoint: {llm_resource.endpoint}")
            logger.info(f"[reflection_lm] Prompt length: {len(prompt)} chars")
            logger.debug(f"[reflection_lm] Prompt preview: {prompt[:500]}...")
            
            try:
                # Route through the LLMProxy for tracing
                # The proxy already has the API key and model config from gepa_runner.py
                logger.info(f"[reflection_lm] Calling via proxy: model={llm_resource.model} endpoint={llm_resource.endpoint}")
                
                # Include detailed tracing metadata for reflection traces
                # Uses run_tracing_config which includes the session_id
                # Note: call-type tagging ("reflection-call") is handled by @gepa_module.reflection decorator
                import json
                gepa_tags = [
                    f"gen-{tracing_context.generation}",
                ]
                tracing_metadata = run_tracing_config.to_detailed_metadata(
                    phase="reflection",
                    extra_tags=gepa_tags,
                )
                
                # Pass metadata via headers for proxy-side tracing context
                # Note: call-type is handled automatically by the @gepa_module.reflection decorator
                extra_headers = {
                    "x-mantis-session-id": tracing_metadata.get("session_id", ""),
                    "x-mantis-environment": tracing_metadata.get("environment", ""),
                    "x-mantis-tags": json.dumps(tracing_metadata.get("tags", [])),
                }
                
                response = litellm.completion(
                    model=llm_resource.model,  # Model name as registered in proxy
                    messages=[{"role": "user", "content": prompt}],
                    api_base=llm_resource.endpoint,  # Proxy URL
                    api_key="dummy",  # Proxy handles the real API key
                    temperature=0.7,
                    metadata=tracing_metadata,  # Pass environment/tags/session_id for OTEL
                    extra_headers=extra_headers,  # Pass to proxy for server-side span attributes
                )
                content = response.choices[0].message.content or ""
                logger.info(f"[reflection_lm] SUCCESS - Response length: {len(content)} chars")
                logger.info(f"[reflection_lm] Response preview: {content[:300]}...")
                
                # After reflection completes, increment generation for next eval round
                tracing_context.next_generation()
                tracing_context.set_phase("train-eval")
                
                return content
            except Exception as e:
                logger.error(f"[reflection_lm] FAILED: {e}", exc_info=True)
                # Return fallback instead of crashing - allows optimization to continue
                logger.warning(f"[reflection_lm] Returning fallback (seed prompt)")
                
                # Still increment generation even on failure
                tracing_context.next_generation()
                tracing_context.set_phase("train-eval")
                
                return f"```\n{initial_candidate.get(target_resource_name, '')}\n```"

        # Prepare Training Data
        gepa_train_data: List[MantisdkDataInst] = []
        for i, item in enumerate(train_dataset):
            item_id = str(i)
            if isinstance(item, dict):
                gepa_train_data.append({"input": item, "id": item.get("id", item_id)})
            else:
                gepa_train_data.append({"input": {"task": item}, "id": item_id})

        # Prepare Validation Data (use train if not provided)
        gepa_val_data: Optional[List[MantisdkDataInst]] = None
        if val_dataset is not None:
            gepa_val_data = []
            for i, item in enumerate(val_dataset):
                item_id = f"val-{i}"
                if isinstance(item, dict):
                    gepa_val_data.append({"input": item, "id": item.get("id", item_id)})
                else:
                    gepa_val_data.append({"input": {"task": item}, "id": item_id})

        logger.info(f"Starting GEPA optimization with {len(gepa_train_data)} training samples")
        if gepa_val_data:
            logger.info(f"Using {len(gepa_val_data)} validation samples")
        
        # Run Optimization in Executor (GEPA is synchronous)
        def run_optimization():
            return optimize(
                trainset=gepa_train_data,
                valset=gepa_val_data,
                adapter=gepa_adapter,
                seed_candidate=initial_candidate,
                reflection_lm=reflection_lm,
                max_metric_calls=self.max_metric_calls,
                reflection_minibatch_size=self.reflection_minibatch_size,
                **self.gepa_kwargs,
            )

        try:
            result = await loop.run_in_executor(None, run_optimization)
            
            logger.info(f"GEPA optimization finished. Total metric calls: {result.total_metric_calls}")
            logger.info(f"Best candidate index: {result.best_idx}, score: {result.val_aggregate_scores[result.best_idx]}")
            logger.info(f"Best candidate: {result.best_candidate}")

            self._best_candidate = result.best_candidate
            self._best_score = result.val_aggregate_scores[result.best_idx]
            self._full_result = result  # Store full result for candidate history

            # Update the store with the best candidate
            final_resources = {}
            for key, value in result.best_candidate.items():
                original = initial_resources.get(key)
                engine = original.engine if isinstance(original, PromptTemplate) else "f-string"
                final_resources[key] = PromptTemplate(template=value, engine=engine)
                 
            await store.update_resources("gepa-final", final_resources)
            logger.info("Updated store with best candidate resources (version: 'gepa-final').")

            # Send full GEPA result to Insight for experiment tracking
            if hasattr(store, 'complete_job'):
                summary = result.to_dict()
                # Add session_id for trace filtering in the UI
                summary["session_id"] = tracing_context.session_id
                store.complete_job(summary)

        except Exception as e:
            logger.error(f"GEPA optimization failed: {e}")
            raise
