"""Evaluation loop engine for continuous learning."""
from __future__ import annotations

from dataclasses import asdict
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from logos_adk.agent import Agent
    from logos_adk.testing_framework import GoldenExample, JudgeCase, JudgeResult, LLMJudge

from .models import (
    ChampionChallengerReport,
    JudgeDisagreement,
    LearningScorecard,
    LiveChampionChallengerMetrics,
    OfflineChampionChallengerMetrics,
    ProductionRegressionDataset,
    ScoreWindow,
)
from .utils import (
    _estimate_cost_usd,
    _human_quality_score,
    _mean,
    _parse_ts,
    _quality_rows_summary,
)


class EvaluationLoopEngine:
    """Evaluation and comparison helpers built on captured production runs."""

    def __init__(self, store: Any, *, reward_scorer: Any | None = None) -> None:
        """Initialize evaluation engine.
        
        Args:
            store: LearningStore instance
            reward_scorer: Optional RewardScorer for reward-based scoring
        """
        from logos_adk.learning.dataset import DatasetBuilder
        from logos_adk.learning.improvement import RewardModel
        
        self.store = store
        self.datasets = DatasetBuilder(store)
        self.reward_scorer = reward_scorer or RewardModel()

    def resolve_reward_scorer(
        self,
        *,
        mode: str | None = None,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 500,
    ) -> Any:
        """Resolve reward scorer based on mode."""
        from logos_adk.learning.improvement import build_reward_scorer
        
        if mode is None:
            return self.reward_scorer
        return build_reward_scorer(
            mode,
            store=self.store,
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            fallback_scorer=self.reward_scorer,
        )

    def build_regression_dataset(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
    ) -> ProductionRegressionDataset:
        """Build regression dataset from production runs."""
        from logos_adk.testing_framework import GoldenExample
        
        examples = self.datasets.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            include_unreviewed=False,
        )
        golden: list[GoldenExample] = []
        source_run_ids: list[str] = []
        for example in examples:
            expected = example.corrected_output
            if not expected:
                for feedback in example.feedback:
                    if feedback.expected_output:
                        expected = feedback.expected_output
                        break
            if not expected:
                continue
            is_bad_case = bool(example.failure_types)
            if not is_bad_case:
                is_bad_case = any(item.kind in {"thumbs_down", "edited_answer"} for item in example.feedback)
            if not is_bad_case:
                is_bad_case = any(
                    item.kind in {"escalation", "human_takeover"} and bool(item.value)
                    for item in example.outcomes
                )
            if not is_bad_case:
                continue
            golden.append(
                GoldenExample(
                    name=example.run_id,
                    prompt=example.input_text,
                    expected=expected,
                    session_id=example.session_id,
                    mode="exact",
                )
            )
            source_run_ids.append(example.run_id)
        return ProductionRegressionDataset(
            workspace_id=workspace_id,
            agent_name=agent_name,
            examples=golden,
            source_run_ids=source_run_ids,
        )

    def build_regression_suite(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
    ) -> dict[str, Any]:
        """Build regression suite from production runs."""
        dataset = self.build_regression_dataset(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
        )
        return {
            "examples": [asdict(item) for item in dataset.examples],
            "source_run_ids": dataset.source_run_ids,
            "workspace_id": dataset.workspace_id,
            "agent_name": dataset.agent_name,
            "generated_at": dataset.generated_at,
        }

    def build_judge_cases_from_prod(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
    ) -> list[JudgeCase]:
        """Build judge cases from production runs."""
        from logos_adk.testing_framework import JudgeCase
        
        source_examples = self.datasets.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            include_unreviewed=False,
        )
        cases: list[JudgeCase] = []
        for example in source_examples:
            reference = example.corrected_output
            if not reference:
                for feedback in example.feedback:
                    if feedback.expected_output:
                        reference = feedback.expected_output
                        break
            if not reference:
                continue
            criteria = "Answer should be correct, aligned with the human-corrected reference, and grounded."
            if example.failure_types:
                criteria += " Avoid failure types: " + ", ".join(sorted(example.failure_types))
            cases.append(
                JudgeCase(
                    name=example.run_id,
                    prompt=example.input_text,
                    criteria=criteria,
                    reference=reference,
                    session_id=example.session_id,
                )
            )
        return cases

    def track_judge_disagreements(
        self,
        judge_results: list[JudgeResult],
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
        threshold: float = 0.35,
    ) -> list[JudgeDisagreement]:
        """Track disagreements between judge and human feedback."""
        examples = self.datasets.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            include_unreviewed=False,
        )
        examples_by_run = {item.run_id: item for item in examples}
        disagreements: list[JudgeDisagreement] = []
        for result in judge_results:
            example = examples_by_run.get(result.name)
            if example is None:
                continue
            human_score = _human_quality_score(example)
            delta = abs(float(result.score) - human_score)
            human_passed = human_score >= 0.6
            if delta < threshold and result.passed == human_passed:
                continue
            disagreements.append(
                JudgeDisagreement(
                    run_id=example.run_id,
                    trace_id=example.trace_id,
                    judge_score=round(float(result.score), 4),
                    human_score=human_score,
                    delta=round(delta, 4),
                    judge_passed=result.passed,
                    human_passed=human_passed,
                    prompt_version=example.prompt_version,
                    model=example.model,
                    failure_types=list(example.failure_types),
                )
            )
        return disagreements

    async def run_regression_snapshot(
        self,
        agent: Agent,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
        system_prompt_override: str | None = None,
    ) -> dict[str, Any]:
        """Run regression snapshot evaluation."""
        dataset = self.build_regression_dataset(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
        )
        if system_prompt_override is None:
            from logos_adk.testing_framework import run_golden_dataset
            results = await run_golden_dataset(agent, dataset.examples)
            rows = [result.__dict__ for result in results]
        else:
            rows = []
            for example in dataset.examples:
                response = await agent.run(
                    example.prompt,
                    session_id=example.session_id,
                    system_prompt_override=system_prompt_override,
                )
                actual = response.content
                passed = example.expected in actual if example.mode == "contains" else actual == example.expected
                rows.append(
                    {
                        "name": example.name,
                        "passed": passed,
                        "actual": actual,
                        "expected": example.expected,
                        "mode": example.mode,
                    }
                )
        return {
            "kind": "regression",
            "summary": _quality_rows_summary(rows),
            "rows": rows,
            "reports": {},
        }

    async def run_judge_snapshot(
        self,
        agent: Agent,
        *,
        judge: LLMJudge,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
        system_prompt_override: str | None = None,
    ) -> dict[str, Any]:
        """Run judge evaluation snapshot."""
        cases = self.build_judge_cases_from_prod(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
        )
        if system_prompt_override is None:
            from logos_adk.testing_framework import run_llm_judge_eval
            results = await run_llm_judge_eval(agent, cases, judge)
        else:
            results = []
            for case in cases:
                answer = await agent.run(
                    case.prompt,
                    session_id=case.session_id,
                    system_prompt_override=system_prompt_override,
                )
                score, rationale = await judge.evaluate(
                    prompt=case.prompt,
                    answer=answer.content,
                    criteria=case.criteria,
                    reference=case.reference,
                )
                results.append(
                    JudgeResult(
                        name=case.name,
                        score=score,
                        passed=score >= judge.pass_threshold,
                        rationale=rationale,
                        answer=answer.content,
                    )
                )
        rows = [result.__dict__ for result in results]
        disagreements = self.track_judge_disagreements(
            results,
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
        )
        summary = _quality_rows_summary(rows)
        summary["judge_disagreements"] = len(disagreements)
        return {
            "kind": "judge-eval",
            "summary": summary,
            "rows": rows,
            "reports": {},
            "disagreements": [asdict(item) for item in disagreements],
        }

    def build_scorecard(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        window_days: int = 7,
        end_at: datetime | None = None,
        limit: int = 5000,
    ) -> LearningScorecard:
        """Build learning scorecard for time window."""
        anchor = end_at or datetime.now(UTC)
        current_start = anchor - timedelta(days=window_days)
        previous_start = current_start - timedelta(days=window_days)
        runs = self.store.list_runs(workspace_id=workspace_id, agent_name=agent_name, limit=limit)
        examples = self.datasets.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            include_unreviewed=True,
        )
        examples_by_run = {item.run_id: item for item in examples}

        def metrics_for_window(start: datetime, end: datetime) -> tuple[dict[str, float], dict[str, int]]:
            window_runs = [
                run
                for run in runs
                if (timestamp := _parse_ts(run.created_at)) is not None and start <= timestamp < end
            ]
            window_examples = [examples_by_run[run.id] for run in window_runs if run.id in examples_by_run]
            quality_values = [
                _human_quality_score(example)
                for example in window_examples
                if example.feedback or example.outcomes
            ]
            cost_values = [_estimate_cost_usd(run.model, run.usage) for run in window_runs]
            latency_values = [
                float(run.metadata.get("latency_ms", 0.0) or 0.0)
                for run in window_runs
                if isinstance(run.metadata, dict) and run.metadata.get("latency_ms") is not None
            ]
            reviewed_examples = [example for example in window_examples if example.feedback or example.outcomes]
            bad_retrieval_count = sum(
                1 for example in reviewed_examples if "bad_retrieval" in example.failure_types
            )
            retrieval_relevance = (
                max(0.0, 1.0 - (bad_retrieval_count / len(reviewed_examples)))
                if reviewed_examples
                else 0.5
            )
            tool_runs = [run for run in window_runs if run.tool_calls]
            tool_failures = sum(
                1
                for run in tool_runs
                if "wrong_tool" in (examples_by_run.get(run.id).failure_types if run.id in examples_by_run else [])
            )
            tool_success = (
                max(0.0, 1.0 - (tool_failures / len(tool_runs)))
                if tool_runs
                else 0.5
            )
            return (
                {
                    "quality": round(_mean(quality_values, default=0.5), 4),
                    "cost": round(_mean(cost_values, default=0.0), 8),
                    "latency": round(_mean(latency_values, default=0.0), 4),
                    "retrieval_relevance": round(retrieval_relevance, 4),
                    "tool_success": round(tool_success, 4),
                },
                {
                    "runs": len(window_runs),
                    "reviewed_runs": len(reviewed_examples),
                    "tool_runs": len(tool_runs),
                },
            )

        current_metrics, current_counts = metrics_for_window(current_start, anchor)
        previous_metrics, previous_counts = metrics_for_window(previous_start, current_start)

        def window(name: str) -> ScoreWindow:
            current_value = current_metrics[name]
            previous_value = previous_metrics[name]
            return ScoreWindow(
                current=current_value,
                previous=previous_value,
                delta=round(current_value - previous_value, 4),
            )

        return LearningScorecard(
            workspace_id=workspace_id,
            agent_name=agent_name,
            window_days=window_days,
            current_window_start=current_start.isoformat(),
            current_window_end=anchor.isoformat(),
            previous_window_start=previous_start.isoformat(),
            previous_window_end=current_start.isoformat(),
            quality=window("quality"),
            cost=window("cost"),
            latency=window("latency"),
            retrieval_relevance=window("retrieval_relevance"),
            tool_success=window("tool_success"),
            counts={
                "current": current_counts,
                "previous": previous_counts,
            },
        )

    async def compare_champion_challenger(
        self,
        agent: Agent,
        *,
        challenger_prompt: str,
        champion_prompt: str | None = None,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        examples: list[GoldenExample] | None = None,
        champion_filter: dict[str, Any] | None = None,
        challenger_filter: dict[str, Any] | None = None,
        limit: int = 100,
        reward_scorer_mode: str | None = None,
        reward_training_limit: int = 500,
    ) -> ChampionChallengerReport:
        """Compare champion vs challenger configurations."""
        scorer = self.resolve_reward_scorer(
            mode=reward_scorer_mode,
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=reward_training_limit,
        )
        scorer_mode = reward_scorer_mode or "heuristic"
        champion_content = champion_prompt or getattr(getattr(agent, "_config", None), "system_prompt", "") or ""
        dataset = examples
        if dataset is None:
            dataset = self.build_regression_dataset(
                workspace_id=workspace_id,
                agent_name=agent_name,
                limit=limit,
            ).examples
        champion_offline = await self._offline_variant_metrics(
            agent,
            label="champion",
            prompt_content=champion_content,
            examples=dataset,
        )
        challenger_offline = await self._offline_variant_metrics(
            agent,
            label="challenger",
            prompt_content=challenger_prompt,
            examples=dataset,
        )
        champion_live = self._live_variant_metrics(
            label="champion",
            workspace_id=workspace_id,
            agent_name=agent_name,
            filters=champion_filter,
            limit=max(limit * 20, 500),
            reward_scorer=scorer,
        )
        challenger_live = self._live_variant_metrics(
            label="challenger",
            workspace_id=workspace_id,
            agent_name=agent_name,
            filters=challenger_filter,
            limit=max(limit * 20, 500),
            reward_scorer=scorer,
        )
        champion_score = champion_offline.pass_rate + champion_live.reward_score
        challenger_score = challenger_offline.pass_rate + challenger_live.reward_score
        recommendation = "challenger" if challenger_score >= champion_score else "champion"
        return ChampionChallengerReport(
            champion={
                "offline": asdict(champion_offline),
                "live": asdict(champion_live),
            },
            challenger={
                "offline": asdict(challenger_offline),
                "live": asdict(challenger_live),
            },
            recommendation=recommendation,
            reward_scorer_mode=scorer_mode,
        )

    async def _offline_variant_metrics(
        self,
        agent: Agent,
        *,
        label: str,
        prompt_content: str,
        examples: list[GoldenExample],
    ) -> OfflineChampionChallengerMetrics:
        """Calculate offline metrics for variant."""
        if not examples:
            return OfflineChampionChallengerMetrics(label=label, pass_rate=0.0, sample_size=0)
        results: list[float] = []
        for example in examples:
            response = await agent.run(
                example.prompt,
                session_id=example.session_id,
                system_prompt_override=prompt_content,
            )
            passed = example.expected in response.content if example.mode == "contains" else response.content == example.expected
            results.append(1.0 if passed else 0.0)
        return OfflineChampionChallengerMetrics(
            label=label,
            pass_rate=round(_mean(results, default=0.0), 4),
            sample_size=len(examples),
        )

    def _live_variant_metrics(
        self,
        *,
        label: str,
        workspace_id: str | None,
        agent_name: str | None,
        filters: dict[str, Any] | None,
        limit: int,
        reward_scorer: Any | None = None,
    ) -> LiveChampionChallengerMetrics:
        """Calculate live traffic metrics for variant."""
        runs = self.store.list_runs(workspace_id=workspace_id, agent_name=agent_name, limit=limit)
        matched = [run for run in runs if self._matches_filters(run, filters or {})]
        examples = {
            item.run_id: item
            for item in self.datasets.build(
                workspace_id=workspace_id,
                agent_name=agent_name,
                limit=limit,
                include_unreviewed=True,
            )
        }
        quality_values: list[float] = []
        cost_values: list[float] = []
        latency_values: list[float] = []
        conversion_values: list[float] = []
        reward_values: list[float] = []
        scorer = reward_scorer or self.reward_scorer
        for run in matched:
            example = examples.get(run.id)
            quality = _human_quality_score(example) if example is not None else 0.5
            cost = _estimate_cost_usd(run.model, run.usage)
            latency = 0.0
            if isinstance(run.metadata, dict):
                latency = float(run.metadata.get("latency_ms", 0.0) or 0.0)
            conversions = 0.0
            if example is not None:
                for outcome in example.outcomes:
                    if outcome.kind == "conversion" and bool(outcome.value):
                        conversions = 1.0
                        break
            if example is not None:
                reward = scorer.score_example(example).score
            else:
                reward = scorer.score_run(
                    run,
                    feedback=self.store.list_feedback(run_id=run.id),
                    outcomes=self.store.list_outcomes(run_id=run.id),
                ).score
            quality_values.append(quality)
            cost_values.append(cost)
            latency_values.append(latency)
            conversion_values.append(conversions)
            reward_values.append(reward)
        return LiveChampionChallengerMetrics(
            label=label,
            runs=len(matched),
            quality_score=round(_mean(quality_values, default=0.0), 4),
            avg_cost_usd=round(_mean(cost_values, default=0.0), 8),
            avg_latency_ms=round(_mean(latency_values, default=0.0), 4),
            conversion_rate=round(_mean(conversion_values, default=0.0), 4),
            reward_score=round(_mean(reward_values, default=0.0), 4),
        )

    def _matches_filters(self, run: Any, filters: dict[str, Any]) -> bool:
        """Check if run matches filters."""
        metadata = getattr(run, "metadata", {}) or {}
        for key, expected in filters.items():
            if key in {"prompt_version", "model", "config_fingerprint"}:
                if getattr(run, key, None) != expected:
                    return False
            elif metadata.get(key) != expected:
                return False
        return True


__all__ = ["EvaluationLoopEngine"]
