"""Prompt optimization, recommendation, and promotion guard logic."""
from __future__ import annotations

from dataclasses import asdict, dataclass, field
from fnmatch import fnmatch
from typing import Any
from uuid import uuid4

from logos_adk.agent import Agent
from logos_adk.cost import DEFAULT_RATES
from logos_adk.testing_framework import GoldenExample

from .dataset import DatasetBuilder
from .prompt_registry import PromptRecommendation, PromptRegistry
from .store import LearningStore


@dataclass
class PromptCandidateScore:
    """Offline prompt candidate score."""

    prompt_name: str
    content: str
    golden_pass_rate: float = 0.0
    feedback_alignment_rate: float = 0.0
    combined_score: float = 0.0
    sample_size: int = 0


@dataclass
class PromptExperimentReport:
    """Aggregated live experiment metrics."""

    experiment_id: str
    prompt_name: str
    variants: dict[str, dict[str, Any]] = field(default_factory=dict)


@dataclass
class PromptPolicyGuardResult:
    """Promotion decision after cost/latency/safety checks."""

    allowed: bool
    reasons: list[str] = field(default_factory=list)
    baseline_metrics: dict[str, Any] = field(default_factory=dict)
    candidate_metrics: dict[str, Any] = field(default_factory=dict)


def _estimate_cost_usd(model: str | None, usage: dict[str, Any] | None) -> float:
    if not model or not usage:
        return 0.0
    input_tokens = int(usage.get("input_tokens", 0) or 0)
    output_tokens = int(usage.get("output_tokens", 0) or 0)
    rates = None
    if model in DEFAULT_RATES:
        rates = DEFAULT_RATES[model]
    else:
        for pattern, candidate in DEFAULT_RATES.items():
            if "*" in pattern and fnmatch(model, pattern):
                rates = candidate
                break
    if rates is None:
        return 0.0
    return round((input_tokens * rates[0] + output_tokens * rates[1]) / 1_000_000, 8)


class RecommendationEngine:
    """Generate prompt patch recommendations from recurring failures."""

    FAILURE_PATCHES: dict[str, str] = {
        "hallucination": "Do not invent facts. If the answer is not grounded in available context, say that directly.",
        "style_miss": "Prefer concise, factual, and non-hype language tailored to the requested tone.",
        "policy_violation": "Avoid unsafe, non-compliant, or policy-breaking claims. Refuse or redirect when needed.",
        "wrong_tool": "Only call tools when they materially improve correctness or completeness.",
        "bad_retrieval": "Base the answer on retrieved context first and acknowledge missing evidence.",
        "latency": "Prefer direct answers and limit unnecessary tool/toolchain steps.",
        "too_expensive": "Prefer shorter outputs and avoid unnecessary long reasoning or redundant tool use.",
    }

    def __init__(self, store: LearningStore, registry: PromptRegistry) -> None:
        self.store = store
        self.registry = registry

    def suggest_patch(
        self,
        *,
        prompt_name: str,
        workspace_id: str | None = None,
        limit: int = 100,
    ) -> PromptRecommendation:
        builder = DatasetBuilder(self.store)
        examples = builder.build(workspace_id=workspace_id, agent_name=prompt_name, limit=limit, include_unreviewed=False)
        failure_counts: dict[str, int] = {}
        corrected_samples: list[str] = []
        based_on_version: int | None = None
        for example in examples:
            prompt_version = example.prompt_version
            if based_on_version is None and prompt_version.startswith("v") and prompt_version[1:].isdigit():
                based_on_version = int(prompt_version[1:])
            for failure in example.failure_types:
                failure_counts[failure] = failure_counts.get(failure, 0) + 1
            if example.corrected_output:
                corrected_samples.append(example.corrected_output)
        if not failure_counts:
            failure_counts = {"no_dominant_failure": 1}
        ordered_failures = sorted(
            ((name, count) for name, count in failure_counts.items()),
            key=lambda item: item[1],
            reverse=True,
        )
        patch_lines = [
            self.FAILURE_PATCHES[name]
            for name, _ in ordered_failures
            if name in self.FAILURE_PATCHES
        ]
        if corrected_samples:
            patch_lines.append(
                "Use successful human edits as style exemplars: "
                + corrected_samples[0][:220]
            )
        latest_versions = self.registry.list_versions(prompt_name)
        current_content = latest_versions[0].content if latest_versions else ""
        proposed_patch = "\n".join(f"+ {line}" for line in patch_lines) if patch_lines else "+ No patch generated"
        recommendation = PromptRecommendation(
            id=uuid4().hex,
            prompt_name=prompt_name,
            kind="prompt_patch",
            title=f"Patch suggestion for {prompt_name}",
            summary="Recurring failures suggest updating the operating instructions.",
            proposed_patch=f"--- current\n{current_content}\n+++ proposed_additions\n{proposed_patch}",
            based_on_version=based_on_version,
            failure_counts=failure_counts,
            metadata={"workspace_id": workspace_id, "sample_size": len(examples)},
        )
        self.registry.save_recommendation(recommendation)
        return recommendation


class PromptPolicyGuard:
    """Block promotion when live metrics regress beyond thresholds."""

    def __init__(self, store: LearningStore) -> None:
        self.store = store

    def evaluate(
        self,
        *,
        baseline_metrics: dict[str, Any],
        candidate_metrics: dict[str, Any],
        max_latency_increase_ratio: float = 1.15,
        max_cost_increase_ratio: float = 1.15,
    ) -> PromptPolicyGuardResult:
        reasons: list[str] = []
        baseline_latency = float(baseline_metrics.get("avg_latency_ms", 0.0) or 0.0)
        candidate_latency = float(candidate_metrics.get("avg_latency_ms", 0.0) or 0.0)
        if baseline_latency > 0 and candidate_latency > baseline_latency * max_latency_increase_ratio:
            reasons.append("latency_regression")
        baseline_cost = float(baseline_metrics.get("avg_estimated_cost_usd", 0.0) or 0.0)
        candidate_cost = float(candidate_metrics.get("avg_estimated_cost_usd", 0.0) or 0.0)
        if baseline_cost > 0 and candidate_cost > baseline_cost * max_cost_increase_ratio:
            reasons.append("cost_regression")
        baseline_safety = float(baseline_metrics.get("safety_violation_rate", 0.0) or 0.0)
        candidate_safety = float(candidate_metrics.get("safety_violation_rate", 0.0) or 0.0)
        if candidate_safety > baseline_safety:
            reasons.append("safety_regression")
        return PromptPolicyGuardResult(
            allowed=not reasons,
            reasons=reasons,
            baseline_metrics=baseline_metrics,
            candidate_metrics=candidate_metrics,
        )


class PromptOptimizer:
    """Run offline prompt candidates against golden and feedback-derived datasets."""

    def __init__(self, store: LearningStore, registry: PromptRegistry) -> None:
        self.store = store
        self.registry = registry
        self.recommendations = RecommendationEngine(store, registry)
        self.guard = PromptPolicyGuard(store)

    async def evaluate_candidates(
        self,
        agent: Agent,
        *,
        prompt_name: str,
        candidates: list[str],
        golden_examples: list[GoldenExample] | None = None,
        workspace_id: str | None = None,
        feedback_limit: int = 100,
    ) -> list[PromptCandidateScore]:
        derived_examples = self._feedback_examples(prompt_name, workspace_id=workspace_id, limit=feedback_limit)
        scores: list[PromptCandidateScore] = []
        for content in candidates:
            golden_score = await self._golden_score(agent, content, golden_examples or [])
            feedback_score = await self._feedback_score(agent, content, derived_examples)
            combined = round((golden_score * 0.7) + (feedback_score * 0.3), 4)
            scores.append(
                PromptCandidateScore(
                    prompt_name=prompt_name,
                    content=content,
                    golden_pass_rate=golden_score,
                    feedback_alignment_rate=feedback_score,
                    combined_score=combined,
                    sample_size=len(golden_examples or []) + len(derived_examples),
                )
            )
        return sorted(scores, key=lambda item: item.combined_score, reverse=True)

    async def create_optimized_version(
        self,
        agent: Agent,
        *,
        prompt_name: str,
        candidates: list[str],
        golden_examples: list[GoldenExample] | None = None,
        workspace_id: str | None = None,
        feedback_limit: int = 100,
        tags: list[str] | None = None,
    ) -> dict[str, Any]:
        scores = await self.evaluate_candidates(
            agent,
            prompt_name=prompt_name,
            candidates=candidates,
            golden_examples=golden_examples,
            workspace_id=workspace_id,
            feedback_limit=feedback_limit,
        )
        if not scores:
            raise ValueError("No prompt candidates provided")
        winner = scores[0]
        created = self.registry.create_version(
            prompt_name,
            content=winner.content,
            tags=list(tags or []) + ["optimized"],
            notes=f"combined_score={winner.combined_score}",
        )
        return {
            "created_version": asdict(created),
            "scores": [asdict(score) for score in scores],
        }

    def build_experiment_report(self, experiment_id: str) -> PromptExperimentReport:
        experiment = self.registry.get_experiment(experiment_id)
        if experiment is None:
            raise ValueError(f"Experiment not found: {experiment_id}")
        runs = self.store.list_runs(workspace_id=experiment.workspace_id, limit=5000)
        relevant_runs = [
            run
            for run in runs
            if run.metadata.get("prompt_name") == experiment.prompt_name
            and run.metadata.get("experiment_id") == experiment.id
        ]
        report = PromptExperimentReport(experiment_id=experiment.id, prompt_name=experiment.prompt_name)
        feedback_by_run: dict[str, list[Any]] = {}
        outcome_by_run: dict[str, list[Any]] = {}
        for run in relevant_runs:
            feedback_by_run[run.id] = self.store.list_feedback(run_id=run.id)
            outcome_by_run[run.id] = self.store.list_outcomes(run_id=run.id)
        for run in relevant_runs:
            version = run.metadata.get("prompt_variant") or f"v{run.prompt_version}"
            bucket = report.variants.setdefault(
                str(version),
                {
                    "runs": 0,
                    "positive_feedback": 0,
                    "negative_feedback": 0,
                    "conversions": 0,
                    "resolutions": 0,
                    "human_takeovers": 0,
                    "ctr_values": [],
                    "reply_rate_values": [],
                    "latency_values": [],
                    "cost_values": [],
                    "safety_violations": 0,
                },
            )
            bucket["runs"] += 1
            latency_ms = run.metadata.get("latency_ms")
            if latency_ms is not None:
                bucket["latency_values"].append(float(latency_ms))
            bucket["cost_values"].append(_estimate_cost_usd(run.model, run.usage))
            for item in feedback_by_run.get(run.id, []):
                if item.kind == "thumbs_up" or (item.kind == "numeric_rating" and (item.rating or 0) >= 4):
                    bucket["positive_feedback"] += 1
                if item.kind == "thumbs_down" or (item.kind == "numeric_rating" and (item.rating or 0) < 4):
                    bucket["negative_feedback"] += 1
                if "policy_violation" in item.failure_types:
                    bucket["safety_violations"] += 1
            for item in outcome_by_run.get(run.id, []):
                if item.kind == "conversion" and bool(item.value):
                    bucket["conversions"] += 1
                if item.kind == "resolution" and bool(item.value):
                    bucket["resolutions"] += 1
                if item.kind == "human_takeover" and bool(item.value):
                    bucket["human_takeovers"] += 1
                if item.kind == "ctr" and isinstance(item.value, (int, float)):
                    bucket["ctr_values"].append(float(item.value))
                if item.kind == "reply_rate" and isinstance(item.value, (int, float)):
                    bucket["reply_rate_values"].append(float(item.value))
        for key, bucket in list(report.variants.items()):
            runs_count = max(1, bucket["runs"])
            positive_rate = bucket["positive_feedback"] / runs_count
            conversion_rate = bucket["conversions"] / runs_count
            resolution_rate = bucket["resolutions"] / runs_count
            takeover_rate = bucket["human_takeovers"] / runs_count
            safety_rate = bucket["safety_violations"] / runs_count
            avg_latency = sum(bucket["latency_values"]) / len(bucket["latency_values"]) if bucket["latency_values"] else 0.0
            avg_cost = sum(bucket["cost_values"]) / len(bucket["cost_values"]) if bucket["cost_values"] else 0.0
            avg_ctr = sum(bucket["ctr_values"]) / len(bucket["ctr_values"]) if bucket["ctr_values"] else 0.0
            avg_reply = sum(bucket["reply_rate_values"]) / len(bucket["reply_rate_values"]) if bucket["reply_rate_values"] else 0.0
            reward = round(
                positive_rate + conversion_rate + resolution_rate + avg_ctr + avg_reply - takeover_rate - safety_rate - (avg_cost * 10),
                4,
            )
            report.variants[key] = {
                "runs": bucket["runs"],
                "positive_feedback_rate": round(positive_rate, 4),
                "conversion_rate": round(conversion_rate, 4),
                "resolution_rate": round(resolution_rate, 4),
                "human_takeover_rate": round(takeover_rate, 4),
                "safety_violation_rate": round(safety_rate, 4),
                "avg_latency_ms": round(avg_latency, 2),
                "avg_estimated_cost_usd": round(avg_cost, 8),
                "avg_ctr": round(avg_ctr, 4),
                "avg_reply_rate": round(avg_reply, 4),
                "reward_score": reward,
            }
        return report

    async def _golden_score(self, agent: Agent, prompt_content: str, examples: list[GoldenExample]) -> float:
        if not examples:
            return 0.0
        results = []
        for example in examples:
            response = await agent.run(
                example.prompt,
                session_id=example.session_id,
                system_prompt_override=prompt_content,
            )
            actual = response.content
            if example.mode == "contains":
                passed = example.expected in actual
            else:
                passed = actual == example.expected
            results.append(1.0 if passed else 0.0)
        return round(sum(results) / len(results), 4)

    async def _feedback_score(self, agent: Agent, prompt_content: str, examples: list[GoldenExample]) -> float:
        if not examples:
            return 0.0
        results = []
        for example in examples:
            response = await agent.run(
                example.prompt,
                session_id=example.session_id,
                system_prompt_override=prompt_content,
            )
            actual = response.content
            if example.mode == "contains":
                passed = example.expected in actual
            else:
                passed = actual == example.expected
            results.append(1.0 if passed else 0.0)
        return round(sum(results) / len(results), 4)

    def _feedback_examples(
        self,
        prompt_name: str,
        *,
        workspace_id: str | None,
        limit: int,
    ) -> list[GoldenExample]:
        builder = DatasetBuilder(self.store)
        examples = builder.build(workspace_id=workspace_id, agent_name=prompt_name, limit=limit, include_unreviewed=False)
        derived: list[GoldenExample] = []
        for example in examples:
            if example.corrected_output:
                derived.append(
                    GoldenExample(
                        name=example.run_id,
                        prompt=example.input_text,
                        expected=example.corrected_output,
                        session_id=example.session_id,
                    )
                )
        return derived


__all__ = [
    "PromptCandidateScore",
    "PromptExperimentReport",
    "PromptPolicyGuardResult",
    "PromptOptimizer",
    "PromptPolicyGuard",
    "RecommendationEngine",
]
