# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

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

import openai

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.lib.api import optimize
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
    """

    # Platform metadata
    name: ClassVar[str] = "gepa"
    display_name: ClassVar[str] = "GEPA"
    description: ClassVar[str] = "Genetic-Pareto Prompt Adaptation"
    requirements: ClassVar[Dict[str, Any]] = {"gpu": False, "min_runners": 1}
    config_schema: ClassVar[Optional[Dict[str, Any]]] = {
        "type": "object",
        "properties": {
            "reflection_model": {
                "type": "string",
                "default": "anthropic/claude-haiku-4-5",
                "description": "LLM model used by the algorithm to analyze agent failures and generate improved prompts. Accepts any LiteLLM model string.",
            },
            "max_metric_calls": {
                "type": "integer",
                "default": 1000,
                "minimum": 1,
                "description": "Maximum number of evaluations (budget)",
            },
            "population_size": {
                "type": "integer",
                "default": 5,
                "minimum": 1,
                "maximum": 50,
                "description": "Size of the population in evolutionary search",
            },
            "reflection_minibatch_size": {
                "type": "integer",
                "default": 3,
                "minimum": 1,
                "description": "Number of examples per reflection batch",
            },
            "rollout_batch_timeout": {
                "type": "number",
                "default": 300.0,
                "minimum": 30.0,
                "description": "Timeout in seconds for waiting for rollout results",
            },
            "candidate_selection_strategy": {
                "type": "string",
                "default": "pareto",
                "enum": ["pareto", "current_best", "epsilon_greedy"],
                "description": "Strategy for selecting the next candidate to optimize",
            },
            "perfect_score": {
                "type": "number",
                "default": 1.0,
                "description": "Score at which an example is considered perfectly solved",
            },
            "seed": {
                "type": "integer",
                "default": 0,
                "description": "Random seed for reproducibility",
            },
            "initial_prompt": {
                "type": "string",
                "description": "Initial system prompt template to optimize (uses {input} placeholder)",
            },
            "val_fraction": {
                "type": "number",
                "default": 0.2,
                "minimum": 0.0,
                "maximum": 0.5,
                "description": "Fraction of dataset held out for validation",
            },
        },
    }

    # Algorithm-owned tracing configuration
    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_model: str = "anthropic/claude-haiku-4-5",
        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_model: LiteLLM model string for the reflection LLM.
            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_model = reflection_model
        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
        self._engine: Optional[Any] = None  # GEPAEngine reference for cancellation

    def request_cancel(self) -> None:
        """Signal GEPA to stop at the next iteration boundary."""
        super().request_cancel()
        if self._engine is not None:
            self._engine.request_stop()
            logger.info("GEPA engine stop requested via cancellation")

    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=False)
    async def run(
        self,
        llm_proxy,  # injected by decorator (may be None)
        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

        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 if available, otherwise direct
        llm_resource = llm_proxy.as_resource() if llm_proxy is not None else None
        reflection_model = self.reflection_model

        if llm_resource is not None:
            logger.info(f"GEPA reflection via LLM proxy: {llm_resource.model}")
        else:
            logger.info(f"GEPA reflection via direct litellm: {reflection_model}")

        # 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

        reflection_failures = 0

        @gepa_module.reflection
        def reflection_lm(prompt: str) -> str:
            """GEPA reflection LM callable -- string in, string out.

            Two modes:
            - With LLM proxy: routes through the proxy for centralized
              key management and tracing.
            - Without proxy: calls litellm directly using API keys from
              the environment.  Tracing still works via tracing.init().
            """
            import litellm

            tracing_context.set_phase("reflection")

            model_name = llm_resource.model if llm_resource else reflection_model
            logger.info(f"[reflection_lm] === REFLECTION (gen-{tracing_context.generation}) ===")
            logger.info(f"[reflection_lm] Model: {model_name}, prompt: {len(prompt)} chars")

            try:
                from mantisdk.tracing.context import tracing_context as _tracing_ctx

                gepa_tags = run_tracing_config.get_detailed_tags(
                    phase="reflection",
                    extra_tags=[f"gen-{tracing_context.generation}"],
                )
                reflection_session = f"{tracing_context.session_id}-reflection-gen{tracing_context.generation}"

                with _tracing_ctx(
                    tags=gepa_tags, environment=run_tracing_config.environment, session_id=reflection_session
                ):
                    if llm_resource is not None:
                        # Route through LLM proxy
                        proxy_endpoint = llm_resource.get_base_url(None, None)
                        response = litellm.completion(
                            model=llm_resource.model,
                            messages=[{"role": "user", "content": prompt}],
                            api_base=proxy_endpoint,
                            api_key="dummy",
                            temperature=0.7,
                            extra_headers={"x-mantis-session-id": reflection_session},
                        )
                    else:
                        # Direct litellm call -- API key from environment
                        response = litellm.completion(
                            model=reflection_model,
                            messages=[{"role": "user", "content": prompt}],
                            temperature=0.7,
                        )

                content = response.choices[0].message.content or ""
                logger.info(f"[reflection_lm] SUCCESS - {len(content)} chars")

                tracing_context.next_generation()
                tracing_context.set_phase("train-eval")
                return content

            except Exception as e:
                nonlocal reflection_failures
                reflection_failures += 1
                logger.error(f"[reflection_lm] FAILED ({reflection_failures} total): {e}", exc_info=True)
                if reflection_failures > 3:
                    raise RuntimeError(
                        f"GEPA reflection failed {reflection_failures} times consecutively. "
                        "Check LLM endpoint configuration."
                    ) from e
                logger.warning(f"[reflection_lm] Returning fallback (seed prompt)")
                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")

        # Create live state callback for real-time visualization
        def on_step_callback(state) -> None:
            """Update live state for real-time visualization after each GEPA iteration."""
            logger.info(f"[on_step_callback] Called for iteration {state.i}")

            if not hasattr(store, "log_state"):
                logger.info(f"[on_step_callback] store has no log_state method")
                return

            logger.info(f"[on_step_callback] calling store.log_state")

            try:
                # Build candidates list with full data
                candidates = []
                for idx, cand in enumerate(state.program_candidates):
                    # Get prompt text from candidate dict
                    prompt_text = ""
                    prompt_key = ""
                    if cand:
                        first_key = next(iter(cand.keys()), None)
                        if first_key:
                            prompt_key = first_key
                            prompt_text = cand.get(first_key, "")

                    # Get score from prog_candidate_val_subscores (dict of val_id -> score)
                    val_scores = (
                        state.prog_candidate_val_subscores[idx] if idx < len(state.prog_candidate_val_subscores) else {}
                    )
                    score = sum(val_scores.values()) / len(val_scores) if val_scores else None

                    # Get parent indices from parent_program_for_candidate
                    parent_list = (
                        state.parent_program_for_candidate[idx] if idx < len(state.parent_program_for_candidate) else []
                    )
                    parent_idx = parent_list[0] if parent_list and parent_list[0] is not None else None

                    candidates.append(
                        {
                            "idx": idx,
                            "prompt": prompt_text[:2000],  # Truncate for sanity
                            "promptKey": prompt_key,
                            "score": score,
                            "parent": parent_idx,
                        }
                    )

                # Build pareto front from program_at_pareto_front_valset
                pareto_front = {}
                if hasattr(state, "program_at_pareto_front_valset"):
                    for instance_id, best_idxs in state.program_at_pareto_front_valset.items():
                        pareto_front[str(instance_id)] = list(best_idxs) if best_idxs else []

                # Find best candidate by computing average validation score
                best_idx = 0
                best_score = 0.0
                for idx, val_scores in enumerate(state.prog_candidate_val_subscores):
                    if val_scores:
                        avg = sum(val_scores.values()) / len(val_scores)
                        if avg > best_score:
                            best_score = avg
                            best_idx = idx

                # Log the live state via public API (F14)
                store.log_state(
                    {
                        "candidates": candidates,
                        "pareto_front": pareto_front,
                        "best_idx": best_idx,
                        "best_score": best_score,
                        "generation": tracing_context.generation,
                        "total_metric_calls": state.total_num_evals,
                        "iteration": state.i,
                    }
                )

                logger.info(
                    f"[on_step_callback] Updated live state: {len(candidates)} candidates, best={best_score:.3f}"
                )

                # Log per-step metrics for the Metrics tab time series
                if hasattr(store, "_run") and store._run:
                    try:
                        store._run.log(
                            {
                                "best_score": best_score,
                                "num_candidates": len(candidates),
                                "metric_calls": state.total_num_evals,
                            },
                            step=max(0, state.i),
                        )
                    except Exception:
                        pass  # Metrics are best-effort; Evolution tab is primary
            except Exception as e:
                logger.warning(f"[on_step_callback] Failed to update live state: {e}")

        # Create a cancellation stopper that checks Algorithm._cancel_event
        cancel_event = self._cancel_event

        class _CancelStopper:
            """Bridges Algorithm.request_cancel() → GEPA engine stop."""

            def __call__(self, gepa_state: Any) -> bool:
                return cancel_event.is_set()

        # Merge with any existing stop_callbacks from gepa_kwargs
        extra_kwargs = dict(self.gepa_kwargs)
        existing_stop_cbs = extra_kwargs.pop("stop_callbacks", None)
        stop_cbs: list = []
        if existing_stop_cbs is not None:
            if isinstance(existing_stop_cbs, (list, tuple)):
                stop_cbs.extend(existing_stop_cbs)
            else:
                stop_cbs.append(existing_stop_cbs)
        stop_cbs.append(_CancelStopper())

        # 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,
                on_step_callback=on_step_callback,
                stop_callbacks=stop_cbs,
                **extra_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

            # Log final summary metrics (these complement the per-step
            # best_score series logged by on_step_callback above)
            if hasattr(store, "_run") and store._run:
                try:
                    store._run.log(
                        {
                            "best_score": self._best_score,
                            "num_candidates": len(result.candidates),
                            "metric_calls": result.total_metric_calls,
                        },
                        step=result.total_metric_calls,
                    )
                    logger.info(f"Logged final GEPA metrics: best_score={self._best_score}")
                except Exception as e:
                    logger.warning(f"Failed to log final metrics: {e}")

            # 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
