"""
Evox context builder: co-evolves search algorithms and solutions.
"""

import asyncio
import concurrent.futures
import hashlib
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
from skydiscover.config import PromptConfig, LLMModelConfig
from skydiscover.llm.llm_pool import LLMPool
from skydiscover.prompt.default.prompt import DefaultContextBuilder
from skydiscover.prompt.templates import TemplateManager
from skydiscover.prompt.evox import formatters as fmt


def run_async_safely(coro):
    """Run an async coroutine, handling nested event loops."""
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(coro)
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(asyncio.run, coro)
            return future.result()

if TYPE_CHECKING:
    from skydiscover.search.base_database import Program

logger = logging.getLogger(__name__)


class EvoxContextBuilder(DefaultContextBuilder):
    """Context builder for search algorithm discovery."""

    def __init__(
        self,
        config: PromptConfig,
        summary_llm_config: List[LLMModelConfig],
        use_llm_stats_insight: Optional[bool] = True,
    ):
        super().__init__(config)
        self.use_llm_stats_insight = use_llm_stats_insight

        evox_templates_dir = Path(__file__).parent / "templates"
        self.template_manager = TemplateManager(custom_template_dir=str(evox_templates_dir))

        self.summary_llm: LLMPool = LLMPool(summary_llm_config)
        logger.info(f"Initialized guide LLM inside EvoxContextBuilder: {summary_llm_config[0].name}")

        self._problem_context_summary_cache: Dict[str, str] = {}

        evox_search_sys_prompt_path = Path(__file__).parent.parent.parent / "search" / "evox" / "config" / "evox_search_sys_prompt.txt"
        with open(evox_search_sys_prompt_path, "r") as f:
            self.relevant_task_description_message = f.read()

        batch_prompt = self.template_manager.get_template("batch_summary_prompt")
        self._batch_sections = self._parse_template_sections(batch_prompt)

    @staticmethod
    def _parse_template_sections(text: str) -> Dict[str, str]:
        """Parse a template with ===SECTION=== markers into a dict of section_name -> content."""
        sections: Dict[str, str] = {}
        current_section = None
        current_lines: List[str] = []
        for line in text.splitlines():
            stripped = line.strip()
            if stripped.startswith("===") and stripped.endswith("===") and len(stripped) > 6:
                if current_section is not None:
                    sections[current_section] = "\n".join(current_lines).strip()
                current_section = stripped[3:-3]
                current_lines = []
            else:
                current_lines.append(line)
        if current_section is not None:
            sections[current_section] = "\n".join(current_lines).strip()
        return sections


    async def _generate_stats_insight_async(self, stats_text: str) -> str:
        """Generate stats insight via LLM."""
        if not stats_text:
            return ""
        user_content = f"Population Statistics:\n\n{stats_text}"
        return await self.summary_llm.generate(
            system_message=self.template_manager.get_template("stats_insight_system_message"),
            messages=[{"role": "user", "content": user_content}]
        )

    async def _generate_problem_context_summary_async(
        self, problem_description: str, evaluator_context: str
    ) -> str:
        """Generate problem context summary via LLM (cached)."""
        cache_key = hashlib.sha256(
            f"{problem_description}|||{evaluator_context}".encode('utf-8')
        ).hexdigest()

        if cache_key in self._problem_context_summary_cache:
            return self._problem_context_summary_cache[cache_key]

        problem_context_input = self.template_manager.get_template("problem_template").format(
            problem_description=problem_description,
            evaluator_context=evaluator_context,
        )

        response = await self.summary_llm.generate(
            system_message=self.template_manager.get_template("problem_context_summary_system_message"),
            messages=[{"role": "user", "content": problem_context_input}]
        )

        self._problem_context_summary_cache[cache_key] = response
        return response

    async def _generate_batch_summaries_async(self, batch_user_message: str) -> str:
        """Generate batch summaries via LLM."""
        return await self.summary_llm.generate(
            system_message=self._batch_sections["SYSTEM"],
            messages=[{"role": "user", "content": batch_user_message}]
        )

    def set_templates(self, system_template: Optional[str] = None, user_template: Optional[str] = None) -> None:
        """Set custom templates to use for this generator."""
        self.system_template_override = system_template
        self.user_template_override = user_template
        logger.info(f"Templates set: user={user_template}")


    def build_prompt(
        self,
        current_program: Union["Program", Dict[str, "Program"]],
        context: Dict[str, Any] = None,
        **kwargs: Any,
    ) -> Dict[str, str]:
        """Build prompt for search algorithm discovery."""
        context = context or {}

        program_metrics = context.get("program_metrics", {})
        inspiration_programs = context.get("inspiration_programs", [])
        previous_programs = context.get("previous_programs", [])
        language = context.get("language", "python")

        user_template_key = self.user_template_override or "search_evolution_full_rewrite_user"
        user_template = self.template_manager.get_template(user_template_key)
        system_message = self._get_system_message()

        search_window_context = fmt.format_search_window_context(context)

        db_stats = context.get("db_stats", {})

        actual_program = list(current_program.values())[0] if isinstance(current_program, dict) else current_program
        horizon = 0
        if actual_program:
            metrics = actual_program.metrics if hasattr(actual_program, 'metrics') else actual_program.get("metrics", {})
            horizon = int(metrics.get("search_window_horizon") or 0)

        if horizon > 0 and db_stats:
            db_stats = fmt.filter_db_stats_by_horizon(db_stats, horizon)

        search_stats = context.get("search_stats") or {}

        stats_insight_data = None
        population_state = ""
        if self.use_llm_stats_insight and db_stats:
            stats_text = fmt.format_population_state(db_stats)
            if stats_text:
                stats_insight_data = stats_text
        else:
            population_state = fmt.format_population_state(db_stats) if db_stats else ""

        problem_description = fmt.format_problem_description(search_stats.get("problem_description"))
        evaluator_context = fmt.format_evaluator_context(search_stats.get("evaluator_context"))

        problem_context_data = self.template_manager.get_template("problem_template").format(
            problem_description=problem_description,
            evaluator_context=evaluator_context,
        )

        all_programs_data = fmt.prepare_search_algorithms_data(inspiration_programs)
        batch_summary_data = None
        if all_programs_data:
            per_program_tpl = self._batch_sections["PER_PROGRAM"]
            combined_content_parts = []
            for prog in all_programs_data:
                part = (
                    f"=== PROGRAM {prog['program_num']} (score={prog['combined_score']:.4f}, improvement={prog['improvement']:.4f}) ===\n"
                    + per_program_tpl.format(
                        task_description=self.relevant_task_description_message,
                        solution=prog['solution'],
                        db_stats_text=prog['db_stats_text'],
                    )
                )
                combined_content_parts.append(part)
            batch_summary_data = self._batch_sections["INSTRUCTIONS"].format(
                num_programs=len(all_programs_data),
                combined_content="\n".join(combined_content_parts),
            )

        # Gather all LLM calls in parallel
        async def gather_llm_calls():
            tasks = []

            if stats_insight_data:
                tasks.append(("stats_insight", self._generate_stats_insight_async(stats_insight_data)))

            has_meaningful_data = (
                problem_description and problem_description.strip() and
                evaluator_context and evaluator_context.strip() and
                not (problem_description.startswith("(No ") and evaluator_context.startswith("(No "))
            )
            if has_meaningful_data:
                tasks.append(("problem_context", self._generate_problem_context_summary_async(problem_description, evaluator_context)))

            if batch_summary_data:
                tasks.append(("batch_summaries", self._generate_batch_summaries_async(batch_summary_data)))

            if not tasks:
                return {}

            results = await asyncio.gather(*[task for _, task in tasks], return_exceptions=True)

            result_dict = {}
            for (name, _), result in zip(tasks, results):
                if isinstance(result, Exception):
                    result_dict[name] = ""
                else:
                    result_dict[name] = result

            return result_dict

        llm_results = run_async_safely(gather_llm_calls())

        if self.use_llm_stats_insight:
            population_state = llm_results.get("stats_insight", population_state)

        problem_context_summary = llm_results.get("problem_context", "")
        problem_template = problem_context_summary.strip() if problem_context_summary else (problem_context_data or "")

        summaries_by_num = fmt.parse_batch_summaries(
            llm_results.get("batch_summaries", ""), all_programs_data
        )

        simplification_threshold = self.config.suggest_simplification_after_chars
        improvement_areas = fmt.identify_search_improvement_areas(
            actual_program, program_metrics, previous_programs, simplification_threshold
        )
        inspiration_section = fmt.format_search_algorithms(
            inspiration_programs, language, summaries_by_num=summaries_by_num
        )

        current_program_section = fmt.format_current_program(
            current_program, language,
            improvement_areas=improvement_areas,
        )

        user_message = user_template.format(
            current_program=current_program_section,
            inspiration_programs=inspiration_section,
            language=language,
            search_window_context=search_window_context,
            population_state=population_state,
            problem_template=problem_template,
            **kwargs,
        )

        return {
            "system": system_message,
            "user": user_message,
        }
