"""Executor: 前向执行，渲染 prompt + 并发调用 LLM + 收集结果。"""

from __future__ import annotations

import re
from concurrent.futures import ThreadPoolExecutor, as_completed

from promptgrad.llm import call_llm
from promptgrad.models import ExecutionResult, TestCase


class Executor:
    """批量执行 prompt + test cases。"""

    def __init__(self, prompt_config: dict, model_config: dict, concurrency: int = 10):
        self.prompt_config = prompt_config
        self.model_config = model_config
        self.concurrency = concurrency

    def render_prompt(self, template: str, variables: dict) -> str:
        """简单 {{var}} 模板渲染。未匹配的变量保留原样。"""

        def replacer(match):
            key = match.group(1).strip()
            return str(variables[key]) if key in variables else match.group(0)

        return re.sub(r"\{\{(\s*\w+\s*)\}\}", replacer, template)

    def _build_messages(self, case_input: dict) -> list[dict]:
        """构建 LLM messages 列表。"""
        messages = []
        if self.prompt_config.get("system"):
            messages.append({"role": "system", "content": self.prompt_config["system"]})
        for ex in self.prompt_config.get("few_shot_examples", []):
            messages.append({"role": "user", "content": ex["user"]})
            messages.append({"role": "assistant", "content": ex["assistant"]})
        user_content = self.render_prompt(self.prompt_config["template"], case_input)
        messages.append({"role": "user", "content": user_content})
        return messages

    def _execute_one(self, case: TestCase) -> ExecutionResult:
        """执行单个 test case。"""
        messages = self._build_messages(case.input)
        result = call_llm(
            model=self.model_config["name"],
            messages=messages,
            temperature=self.model_config.get("temperature", 0.3),
            max_tokens=self.model_config.get("max_tokens", 1000),
        )
        return ExecutionResult(
            case_id=case.id,
            input=case.input,
            output=result["output"],
            latency_ms=result["latency_ms"],
            token_usage=result["token_usage"],
        )

    def run(self, test_cases: list[TestCase]) -> list[ExecutionResult]:
        """并发执行所有 test cases。"""
        results: list[ExecutionResult] = []
        with ThreadPoolExecutor(max_workers=self.concurrency) as pool:
            futures = {pool.submit(self._execute_one, tc): tc for tc in test_cases}
            for future in as_completed(futures):
                results.append(future.result())
        results.sort(key=lambda r: r.case_id)
        return results
