from __future__ import annotations

import asyncio
import logging
import time
from datetime import datetime
from typing import Any

from leshy.attacks.base import Attack
from leshy.buffs.base import Buff
from leshy.core.models import (
    AttackConfig,
    AttackPayload,
    BuffConfig,
    EvalConfig,
    EvalResult,
    TestCaseConfig,
    TestCaseResult,
    TestSuiteConfig,
    TestSuiteResult,
    Verdict,
)
from leshy.core.registry import (
    attack_registry,
    buff_registry,
    evaluator_registry,
    logger_registry,
    target_registry,
)
from leshy.core.session import SessionContext
from leshy.core.tokens import token_tracker
from leshy.evaluators.base import Evaluator
from leshy.loggers.base import Logger
from leshy.targets.base import Target

logger = logging.getLogger(__name__)


class Runner:
    """Orchestrates test suite execution: single-turn and multi-turn."""

    def __init__(
        self,
        config: TestSuiteConfig,
        *,
        tag_filter: str | None = None,
        verbose: bool = False,
        output_file: str | None = None,
    ) -> None:
        self.config = config
        self.tag_filter = tag_filter
        self.verbose = verbose
        self.output_file = output_file

    async def run(self) -> TestSuiteResult:
        token_tracker.reset()

        target = self._resolve_target()
        loggers = self._resolve_loggers()

        suite_result = TestSuiteResult(
            suite_name=self.config.name,
            started_at=datetime.utcnow(),
        )

        await target.setup()
        await self._fan_out(loggers, "on_suite_start", self.config)

        try:
            cases = self._filter_cases()
            for case_config in cases:
                case_result = await self._run_case(case_config, target, loggers)
                suite_result.case_results.append(case_result)
        finally:
            await target.teardown()

        suite_result.completed_at = datetime.utcnow()
        suite_result.total_payloads = sum(c.total_payloads for c in suite_result.case_results)
        suite_result.total_vulnerable = sum(
            c.vulnerable_count for c in suite_result.case_results
        )
        suite_result.total_resilient = sum(c.resilient_count for c in suite_result.case_results)
        suite_result.total_inconclusive = sum(
            c.inconclusive_count for c in suite_result.case_results
        )
        suite_result.total_errors = sum(c.error_count for c in suite_result.case_results)
        suite_result.token_usage = token_tracker.usage

        await self._fan_out(loggers, "on_suite_complete", suite_result)
        return suite_result

    async def _run_case(
        self,
        case_config: TestCaseConfig,
        target: Target,
        loggers: list[Logger],
    ) -> TestCaseResult:
        await self._fan_out(loggers, "on_case_start", case_config)
        start = time.monotonic()

        result = TestCaseResult(
            name=case_config.name,
            attack_names=[ac.name for ac in case_config.attacks],
        )
        all_eval_results: list[EvalResult] = []

        buffs = self._resolve_buffs(case_config.buffs)

        if case_config.multi_turn and hasattr(target, 'config'):
            tmpl = getattr(target.config, 'body_template', '') or ''
            if '{{session_id}}' not in tmpl and '{{thread_id}}' not in tmpl:
                logger.warning(
                    "Multi-turn test '%s' but body_template has no "
                    "{{session_id}} — server won't track conversation context",
                    case_config.name,
                )

        for attack_config in case_config.attacks:
            attack = self._resolve_attack(attack_config)
            evaluators = self._resolve_evaluators(case_config, attack)

            if case_config.multi_turn:
                evals = await self._run_multi_turn(
                    attack, attack_config, target, evaluators, loggers, case_config.max_turns,
                    buffs,
                )
            else:
                concurrency = (
                    case_config.concurrency
                    if case_config.concurrency > 1
                    else self.config.concurrency
                )
                evals = await self._run_single_turn(
                    attack, attack_config, target, evaluators, loggers, buffs,
                    concurrency,
                )
            all_eval_results.extend(evals)

        result.eval_results = all_eval_results
        result.total_payloads = len(all_eval_results)
        result.vulnerable_count = sum(
            1 for r in all_eval_results if r.verdict == Verdict.VULNERABLE
        )
        result.resilient_count = sum(
            1 for r in all_eval_results if r.verdict == Verdict.RESILIENT
        )
        result.inconclusive_count = sum(
            1 for r in all_eval_results if r.verdict == Verdict.INCONCLUSIVE
        )
        result.error_count = sum(1 for r in all_eval_results if r.verdict == Verdict.ERROR)
        result.duration_ms = (time.monotonic() - start) * 1000

        await self._fan_out(loggers, "on_case_complete", result)
        return result

    async def _run_single_turn(
        self,
        attack: Attack,
        attack_config: AttackConfig,
        target: Target,
        evaluators: list[Evaluator],
        loggers: list[Logger],
        buffs: list[Buff],
        concurrency: int = 1,
    ) -> list[EvalResult]:
        payloads = self._resolve_payloads(attack, attack_config)

        if concurrency == 1:
            # Sequential execution (original behavior)
            results: list[EvalResult] = []
            for payload in payloads:
                result = await self._process_payload(
                    payload, target, evaluators, loggers, buffs
                )
                results.append(result)
            return results

        # Parallel execution with semaphore
        semaphore = asyncio.Semaphore(concurrency)

        async def process_with_semaphore(payload: AttackPayload) -> EvalResult:
            async with semaphore:
                return await self._process_payload(
                    payload, target, evaluators, loggers, buffs
                )

        tasks = [process_with_semaphore(p) for p in payloads]
        return list(await asyncio.gather(*tasks))

    async def _process_payload(
        self,
        payload: AttackPayload,
        target: Target,
        evaluators: list[Evaluator],
        loggers: list[Logger],
        buffs: list[Buff],
    ) -> EvalResult:
        """Process a single payload: transform, send, evaluate, log."""
        try:
            transformed = await self._apply_buffs(payload.content, buffs)
            response = await target.send(transformed)
            eval_results = await self._evaluate_all(evaluators, payload, response)
            merged = self._merge_eval_results(eval_results, payload.id)
            merged.payload_content = payload.content
            merged.response_content = response.content
            merged.response_latency_ms = response.latency_ms
            await self._fan_out(loggers, "on_payload_result", payload, response, [merged])
            return merged
        except Exception as e:
            logger.warning("Error processing payload %s: %s", payload.id, e)
            return EvalResult(
                verdict=Verdict.ERROR,
                confidence=0.0,
                reasoning=f"Execution error: {e}",
                evaluator_name="runner",
                payload_id=payload.id,
            )

    async def _run_multi_turn(
        self,
        attack: Attack,
        attack_config: AttackConfig,
        target: Target,
        evaluators: list[Evaluator],
        loggers: list[Logger],
        max_turns: int,
        buffs: list[Buff],
    ) -> list[EvalResult]:
        session = SessionContext()
        results: list[EvalResult] = []

        if attack_config.steps:
            payloads = [
                attack._make_payload(step, turn=i)
                for i, step in enumerate(attack_config.steps)
            ]
            for payload in payloads[:max_turns]:
                try:
                    transformed = await self._apply_buffs(payload.content, buffs)
                    response = await target.send(
                        transformed, thread_id=session.thread_id
                    )
                    session.add_turn(payload, response)
                    eval_results = await self._evaluate_all(evaluators, payload, response)
                    merged = self._merge_eval_results(eval_results, payload.id)
                    merged.payload_content = payload.content
                    merged.response_content = response.content
                    merged.response_latency_ms = response.latency_ms
                    results.append(merged)
                    await self._fan_out(
                        loggers, "on_payload_result", payload, response, [merged]
                    )
                except Exception as e:
                    logger.warning("Error in multi-turn step %d: %s", payload.turn, e)
                    results.append(
                        EvalResult(
                            verdict=Verdict.ERROR,
                            confidence=0.0,
                            reasoning=f"Multi-turn error: {e}",
                            evaluator_name="runner",
                            payload_id=payload.id,
                        )
                    )
        else:
            turn = 0
            async for payload in attack.generate_turns(session):
                if turn >= max_turns:
                    break
                try:
                    transformed = await self._apply_buffs(payload.content, buffs)
                    response = await target.send(
                        transformed, thread_id=session.thread_id
                    )
                    session.add_turn(payload, response)
                    eval_results = await self._evaluate_all(evaluators, payload, response)
                    merged = self._merge_eval_results(eval_results, payload.id)
                    merged.payload_content = payload.content
                    merged.response_content = response.content
                    merged.response_latency_ms = response.latency_ms
                    results.append(merged)
                    await self._fan_out(
                        loggers, "on_payload_result", payload, response, [merged]
                    )
                except Exception as e:
                    logger.warning("Error in multi-turn turn %d: %s", turn, e)
                    results.append(
                        EvalResult(
                            verdict=Verdict.ERROR,
                            confidence=0.0,
                            reasoning=f"Multi-turn error: {e}",
                            evaluator_name="runner",
                            payload_id=payload.id,
                        )
                    )
                turn += 1

        return results

    def _resolve_payloads(
        self, attack: Attack, attack_config: AttackConfig
    ) -> list[AttackPayload]:
        if attack_config.custom_payloads:
            return [
                attack._make_payload(p, turn=0)
                for p in attack_config.custom_payloads
            ]
        return attack.generate_payloads()

    def _resolve_target(self) -> Target:
        cls = target_registry.get(self.config.target.type)
        return cls(config=self.config.target)

    def _resolve_attack(self, attack_config: AttackConfig) -> Attack:
        cls = attack_registry.get(attack_config.name)
        return cls(config=attack_config)

    def _resolve_evaluators(
        self, case_config: TestCaseConfig, attack: Attack
    ) -> list[Evaluator]:
        eval_configs: list[EvalConfig]
        if case_config.evaluators:
            eval_configs = case_config.evaluators
        else:
            eval_configs = attack.default_evaluators

        evaluators: list[Evaluator] = []
        for ec in eval_configs:
            cls = evaluator_registry.get(ec.type)
            params = dict(ec.params)
            if "model" not in params and self.config.llm_model:
                params["model"] = self.config.llm_model
            if "api_key" not in params and self.config.api_key:
                params["api_key"] = self.config.api_key
            evaluators.append(cls(**params))
        return evaluators

    def _resolve_buffs(self, buff_configs: list[BuffConfig]) -> list[Buff]:
        """Resolve buff instances from config."""
        buffs: list[Buff] = []
        for bc in buff_configs:
            cls = buff_registry.get(bc.name)
            params = dict(bc.params)
            if "model" not in params and self.config.llm_model:
                params["model"] = self.config.llm_model
            if "api_key" not in params and self.config.api_key:
                params["api_key"] = self.config.api_key
            buffs.append(cls(**params))
        return buffs

    async def _apply_buffs(self, content: str, buffs: list[Buff]) -> str:
        """Apply buffs in sequence to transform payload content."""
        result = content
        for buff in buffs:
            result = await buff.transform(result)
        return result

    def _resolve_loggers(self) -> list[Logger]:
        result: list[Logger] = []
        for logger_name in self.config.loggers:
            kwargs: dict[str, Any] = {"verbose": self.verbose}
            if logger_name in ("json_file", "html") and self.output_file:
                kwargs["output"] = self.output_file
            cls = logger_registry.get(logger_name)
            result.append(cls(**kwargs))
        return result

    def _filter_cases(self) -> list[TestCaseConfig]:
        if not self.tag_filter:
            return self.config.tests
        return [c for c in self.config.tests if self.tag_filter in c.tags]

    async def _evaluate_all(
        self,
        evaluators: list[Evaluator],
        payload: AttackPayload,
        response: Any,
    ) -> list[EvalResult]:
        tasks = [ev.evaluate(payload, response) for ev in evaluators]
        return list(await asyncio.gather(*tasks))

    @staticmethod
    def _merge_eval_results(results: list[EvalResult], payload_id: str) -> EvalResult:
        """Merge multiple eval results, taking the worst verdict."""
        if not results:
            return EvalResult(
                verdict=Verdict.INCONCLUSIVE,
                confidence=0.0,
                reasoning="No evaluators ran",
                evaluator_name="runner",
                payload_id=payload_id,
            )
        if len(results) == 1:
            return results[0]

        priority = {
            Verdict.VULNERABLE: 0,
            Verdict.ERROR: 1,
            Verdict.INCONCLUSIVE: 2,
            Verdict.RESILIENT: 3,
        }
        worst = min(results, key=lambda r: priority[r.verdict])
        all_reasoning = "; ".join(
            f"[{r.evaluator_name}] {r.reasoning}" for r in results if r.reasoning
        )
        return EvalResult(
            verdict=worst.verdict,
            confidence=worst.confidence,
            reasoning=all_reasoning,
            evaluator_name="merged",
            payload_id=payload_id,
        )

    @staticmethod
    async def _fan_out(loggers: list[Logger], method: str, *args: Any) -> None:
        tasks = [getattr(lgr, method)(*args) for lgr in loggers]
        await asyncio.gather(*tasks, return_exceptions=True)
