"""Closed-loop self-learning orchestration built on existing learning primitives."""
from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any
import json

from logos_adk.agent import Agent

from .dataset import DatasetBuilder
from .evaluation import ContinuousEvalRunner, EvaluationLoopEngine
from .improvement import FailurePattern, PatternMiner, ReflectionMemory, RewardModel, RewardScorer, build_reward_scorer
from .optimizer import PromptCandidateScore, PromptOptimizer
from .prompt_evolution import PromptCandidate, PromptCandidateGenerator, PromptGenerationReport
from .prompt_registry import PromptRegistry
from .store import LearningStore


@dataclass
class SelfLearningLoopPolicy:
    """Thresholds and rollout policy for one self-learning run."""

    min_sample_size: int = 3
    min_improvement: float = 0.05
    max_candidates: int = 5
    feedback_limit: int = 100
    rollout_percentage: int = 20
    environment: str = "prod"
    allow_auto_promote: bool = False
    max_latency_increase_ratio: float = 1.15
    max_cost_increase_ratio: float = 1.15
    auto_rollback_on_regression: bool = True
    rollback_min_pass_rate: float = 0.95
    rollback_min_total_cases: int = 1
    disable_experiment_after: int = 2
    continuous_eval_interval_seconds: float = 3600.0
    continuous_eval_limit: int = 100
    continuous_eval_job_name: str | None = None
    reward_scorer_mode: str = "heuristic"
    reward_training_limit: int = 500


@dataclass
class SelfLearningLoopAction:
    """One deploy/experiment/noop decision from the loop."""

    kind: str
    status: str
    reason: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class SelfLearningLoopRun:
    """Structured result for a self-learning iteration."""

    prompt_name: str
    workspace_id: str | None
    task_type: str | None
    user_segment: str | None
    dry_run: bool
    reward_scorer_mode: str
    baseline_version: int
    baseline_score: float
    winner_score: float
    improvement: float
    historical_reward_score: float
    failure_count: int
    pattern_count: int
    candidate_count: int
    generation_report: dict[str, Any] | None = None
    scores: list[PromptCandidateScore] = field(default_factory=list)
    patterns: list[FailurePattern] = field(default_factory=list)
    action: SelfLearningLoopAction | None = None
    created_version: dict[str, Any] | None = None
    experiment: dict[str, Any] | None = None
    deployment: dict[str, Any] | None = None
    audit_trail: list[str] = field(default_factory=list)


class SelfLearningLoop:
    """Run a minimal self-improvement cycle for prompt evolution."""

    def __init__(
        self,
        store: LearningStore,
        registry: PromptRegistry,
        *,
        prompt_optimizer: PromptOptimizer | None = None,
        evaluation_loop: EvaluationLoopEngine | None = None,
        continuous_eval_runner: ContinuousEvalRunner | None = None,
        pattern_miner: PatternMiner | None = None,
        reflection_memory: ReflectionMemory | None = None,
        candidate_generator: PromptCandidateGenerator | None = None,
        reward_scorer: RewardScorer | None = None,
    ) -> None:
        self.store = store
        self.registry = registry
        self.prompt_optimizer = prompt_optimizer or PromptOptimizer(store, registry)
        self.reward_scorer = reward_scorer or RewardModel()
        self.evaluation_loop = evaluation_loop or EvaluationLoopEngine(store, reward_scorer=self.reward_scorer)
        self.continuous_eval_runner = continuous_eval_runner
        self.pattern_miner = pattern_miner or PatternMiner(store)
        self.reflection_memory = reflection_memory or ReflectionMemory()
        self.candidate_generator = candidate_generator or PromptCandidateGenerator(
            registry,
            store=store,
            reflection_memory=self.reflection_memory,
        )
        self.datasets = DatasetBuilder(store)

    def collect_failures(
        self,
        *,
        prompt_name: str,
        workspace_id: str | None = None,
        limit: int = 100,
    ) -> list[Any]:
        examples = self.datasets.build(
            workspace_id=workspace_id,
            agent_name=prompt_name,
            limit=limit,
            include_unreviewed=False,
        )
        return [item for item in examples if item.failure_types or item.corrected_output]

    def mine_patterns(
        self,
        *,
        prompt_name: str,
        workspace_id: str | None = None,
        limit: int = 100,
    ) -> list[FailurePattern]:
        return self.pattern_miner.mine_patterns(
            workspace_id=workspace_id,
            agent_name=prompt_name,
            limit=limit,
            min_occurrences=1,
        )

    def generate_candidates(
        self,
        *,
        prompt_name: str,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        environment: str = "prod",
        patterns: list[FailurePattern] | None = None,
        max_candidates: int = 5,
    ) -> list[PromptCandidate]:
        return self.candidate_generator.generate(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            environment=environment,
            patterns=patterns,
            max_candidates=max_candidates,
        )

    async def evaluate_candidates(
        self,
        agent: Agent,
        *,
        prompt_name: str,
        candidates: list[str],
        workspace_id: str | None = None,
        feedback_limit: int = 100,
    ) -> list[PromptCandidateScore]:
        regression = self.evaluation_loop.build_regression_dataset(
            workspace_id=workspace_id,
            agent_name=prompt_name,
            limit=feedback_limit,
        )
        return await self.prompt_optimizer.evaluate_candidates(
            agent,
            prompt_name=prompt_name,
            candidates=candidates,
            golden_examples=regression.examples,
            workspace_id=workspace_id,
            feedback_limit=feedback_limit,
        )

    async def run_once(
        self,
        agent: Agent,
        *,
        prompt_name: str,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        policy: SelfLearningLoopPolicy | None = None,
        dry_run: bool = False,
    ) -> SelfLearningLoopRun:
        policy = policy or SelfLearningLoopPolicy()
        resolution = self.registry.resolve(
            prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            environment=policy.environment,
        )
        if resolution is None:
            versions = self.registry.list_versions(prompt_name)
            if not versions:
                raise ValueError(f"Prompt not found: {prompt_name}")
            baseline_version = versions[0].version
            baseline_content = versions[0].content
        else:
            baseline_version = resolution.version
            baseline_content = resolution.content

        failures = self.collect_failures(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            limit=policy.feedback_limit,
        )
        patterns = self.mine_patterns(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            limit=policy.feedback_limit,
        )
        generation_report = self.candidate_generator.generate_report(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            environment=policy.environment,
            patterns=patterns,
            max_candidates=policy.max_candidates,
        )
        candidates = generation_report.candidates
        run = SelfLearningLoopRun(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            dry_run=dry_run,
            reward_scorer_mode=policy.reward_scorer_mode,
            baseline_version=baseline_version,
            baseline_score=0.0,
            winner_score=0.0,
            improvement=0.0,
            historical_reward_score=0.0,
            failure_count=len(failures),
            pattern_count=len(patterns),
            candidate_count=len(candidates),
            generation_report=self._serialize_generation_report(generation_report),
            patterns=patterns,
        )
        reward_scorer = self._resolve_reward_scorer(
            mode=policy.reward_scorer_mode,
            workspace_id=workspace_id,
            prompt_name=prompt_name,
            limit=policy.reward_training_limit,
        )
        run.historical_reward_score = self._historical_reward_score(failures, reward_scorer)
        run.audit_trail.append(f"reward_scorer_mode={policy.reward_scorer_mode}")
        run.audit_trail.append(f"historical_reward={run.historical_reward_score}")

        if len(failures) < policy.min_sample_size:
            run.action = SelfLearningLoopAction(
                kind="noop",
                status="skipped",
                reason="insufficient_failures",
                metadata={"failure_count": len(failures), "min_sample_size": policy.min_sample_size},
            )
            run.audit_trail.append("insufficient_failures")
            return run

        if not candidates:
            run.action = SelfLearningLoopAction(
                kind="noop",
                status="skipped",
                reason="no_candidates_generated",
            )
            run.audit_trail.append("no_candidates_generated")
            return run

        persisted_candidates = self.candidate_generator.persist_candidates(generation_report)
        run.audit_trail.append(f"generated_candidates={len(persisted_candidates)}")

        candidate_map = {candidate.content: candidate for candidate in candidates}
        scores = await self.evaluate_candidates(
            agent,
            prompt_name=prompt_name,
            candidates=[baseline_content, *candidate_map.keys()],
            workspace_id=workspace_id,
            feedback_limit=policy.feedback_limit,
        )
        run.scores = scores
        baseline_score = next((item for item in scores if item.content == baseline_content), None)
        winner_score = scores[0]
        run.baseline_score = baseline_score.combined_score if baseline_score is not None else 0.0
        run.winner_score = winner_score.combined_score
        run.improvement = round(run.winner_score - run.baseline_score, 4)

        if winner_score.content == baseline_content:
            run.action = SelfLearningLoopAction(
                kind="noop",
                status="skipped",
                reason="baseline_still_best",
            )
            run.audit_trail.append("baseline_still_best")
            return run
        if run.improvement < policy.min_improvement:
            run.action = SelfLearningLoopAction(
                kind="noop",
                status="skipped",
                reason="insufficient_lift",
                metadata={"improvement": run.improvement, "min_improvement": policy.min_improvement},
            )
            run.audit_trail.append("insufficient_lift")
            return run

        selected = candidate_map[winner_score.content]
        baseline_metrics = self._live_metrics(
            prompt_name=prompt_name,
            prompt_version=f"v{baseline_version}",
            workspace_id=workspace_id,
        )
        candidate_metrics = self._project_candidate_metrics(
            baseline_metrics=baseline_metrics,
            baseline_content=baseline_content,
            candidate=selected,
            patterns=patterns,
        )
        guard = self.prompt_optimizer.guard.evaluate(
            baseline_metrics=baseline_metrics,
            candidate_metrics=candidate_metrics,
            max_latency_increase_ratio=policy.max_latency_increase_ratio,
            max_cost_increase_ratio=policy.max_cost_increase_ratio,
        )
        run.audit_trail.append(f"winner_strategy={selected.strategy}")
        run.audit_trail.append(f"lift={run.improvement}")
        if not guard.allowed:
            run.audit_trail.append(f"guard_blocked={','.join(guard.reasons)}")
        if dry_run:
            run.action = SelfLearningLoopAction(
                kind="promote" if policy.allow_auto_promote and guard.allowed else "experiment",
                status="dry_run",
                reason="dry_run",
                metadata={"guard_allowed": guard.allowed, "guard_reasons": guard.reasons},
            )
            run.audit_trail.append("dry_run")
            return run

        provenance = {
            "source": "self_learning_loop",
            "strategy": selected.strategy,
            "rationale": selected.rationale,
            "improvement": run.improvement,
            "reward_scorer_mode": policy.reward_scorer_mode,
            "historical_reward_score": run.historical_reward_score,
            "guard_allowed": guard.allowed,
            "patterns": [pattern.signature for pattern in patterns[:5]],
            "generated_candidates": self._serialize_generation_report(generation_report)["candidates"],
        }
        rollback_policy = {
            "prompt_name": prompt_name,
            "workspace_id": workspace_id,
            "task_type": task_type,
            "user_segment": user_segment,
            "environment": policy.environment,
            "auto_rollback_on_regression": policy.auto_rollback_on_regression,
            "min_pass_rate": policy.rollback_min_pass_rate,
            "min_total_cases": policy.rollback_min_total_cases,
            "disable_experiment_after": policy.disable_experiment_after,
            "reward_scorer_mode": policy.reward_scorer_mode,
        }
        created = self.registry.create_version(
            prompt_name,
            content=selected.content,
            tags=["self-learning", selected.strategy],
            notes=json.dumps({**provenance, "rollback_policy": rollback_policy}, sort_keys=True),
            created_by="self_learning_loop",
        )
        run.created_version = asdict(created)

        if policy.allow_auto_promote and guard.allowed:
            deployment = self.registry.promote(
                prompt_name,
                version=created.version,
                environment=policy.environment,
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
                metadata={
                    "source": "self_learning_loop",
                    "baseline_version": baseline_version,
                    "improvement": run.improvement,
                    "rollback_policy": rollback_policy,
                },
            )
            run.deployment = asdict(deployment)
            scheduled_job = self._ensure_continuous_eval_job(
                agent,
                prompt_name=prompt_name,
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
                policy=policy,
                rollback_policy=rollback_policy,
                target_version=created.version,
            )
            run.action = SelfLearningLoopAction(
                kind="promote",
                status="executed",
                reason="auto_promoted",
                metadata={"version": created.version, "continuous_eval_job": scheduled_job},
            )
            run.audit_trail.append(f"promoted=v{created.version}")
            return run

        experiment = self.registry.create_experiment(
            prompt_name,
            baseline_version=baseline_version,
            challenger_versions=[created.version],
            rollout_percentage=policy.rollout_percentage,
            environment=policy.environment,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            metadata={
                "source": "self_learning_loop",
                "baseline_version": baseline_version,
                "improvement": run.improvement,
                "guard_allowed": guard.allowed,
                "guard_reasons": guard.reasons,
                "rollback_policy": rollback_policy,
            },
        )
        run.experiment = asdict(experiment)
        scheduled_job = self._ensure_continuous_eval_job(
            agent,
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            policy=policy,
            rollback_policy=rollback_policy,
            experiment_id=experiment.id,
            target_version=created.version,
        )
        run.action = SelfLearningLoopAction(
            kind="experiment",
            status="executed",
            reason="created_experiment",
            metadata={
                "experiment_id": experiment.id,
                "version": created.version,
                "continuous_eval_job": scheduled_job,
            },
        )
        run.audit_trail.append(f"experiment={experiment.id}")
        return run

    def _live_metrics(
        self,
        *,
        prompt_name: str,
        prompt_version: str,
        workspace_id: str | None,
        limit: int = 200,
    ) -> dict[str, Any]:
        runs = self.store.list_runs(workspace_id=workspace_id, agent_name=prompt_name, limit=limit)
        relevant = []
        for run in runs:
            metadata = run.metadata or {}
            if metadata.get("prompt_name", prompt_name) != prompt_name:
                continue
            if run.prompt_version != prompt_version:
                continue
            relevant.append(run)
        if not relevant:
            return {
                "avg_latency_ms": 0.0,
                "avg_estimated_cost_usd": 0.0,
                "safety_violation_rate": 0.0,
            }
        latencies = [
            float((run.metadata or {}).get("latency_ms", 0.0) or 0.0)
            for run in relevant
            if (run.metadata or {}).get("latency_ms") is not None
        ]
        safety_violations = 0
        for run in relevant:
            for feedback in self.store.list_feedback(run_id=run.id):
                if "policy_violation" in feedback.failure_types:
                    safety_violations += 1
                    break
        return {
            "avg_latency_ms": round(sum(latencies) / len(latencies), 2) if latencies else 0.0,
            "avg_estimated_cost_usd": 0.0,
            "safety_violation_rate": round(safety_violations / max(1, len(relevant)), 4),
        }

    def _project_candidate_metrics(
        self,
        *,
        baseline_metrics: dict[str, Any],
        baseline_content: str,
        candidate: PromptCandidate,
        patterns: list[FailurePattern],
    ) -> dict[str, Any]:
        baseline_latency = float(baseline_metrics.get("avg_latency_ms", 0.0) or 0.0)
        baseline_cost = float(baseline_metrics.get("avg_estimated_cost_usd", 0.0) or 0.0)
        baseline_safety = float(baseline_metrics.get("safety_violation_rate", 0.0) or 0.0)
        length_ratio = len(candidate.content) / max(1, len(baseline_content))
        latency_ratio = min(1.35, max(0.85, length_ratio))
        cost_ratio = min(1.35, max(0.85, length_ratio))
        if candidate.strategy == "compress":
            latency_ratio = min(latency_ratio, 0.9)
            cost_ratio = min(cost_ratio, 0.9)
        candidate_safety = baseline_safety
        targeted = {failure for pattern in patterns for failure in pattern.failure_types}
        if candidate.strategy in {"patch", "rewrite", "retrieval_grounding"} and targeted & {
            "hallucination",
            "bad_retrieval",
            "policy_violation",
        }:
            candidate_safety = max(0.0, baseline_safety - 0.05)
        return {
            "avg_latency_ms": round(baseline_latency * latency_ratio, 2) if baseline_latency else 0.0,
            "avg_estimated_cost_usd": round(baseline_cost * cost_ratio, 8) if baseline_cost else 0.0,
            "safety_violation_rate": round(candidate_safety, 4),
        }

    def _ensure_continuous_eval_job(
        self,
        agent: Agent,
        *,
        prompt_name: str,
        workspace_id: str | None,
        task_type: str | None,
        user_segment: str | None,
        policy: SelfLearningLoopPolicy,
        rollback_policy: dict[str, Any],
        experiment_id: str | None = None,
        target_version: int | None = None,
    ) -> str | None:
        runner = self.continuous_eval_runner
        if runner is None:
            return None
        job_name = policy.continuous_eval_job_name or (
            f"self-learning-{prompt_name}-{workspace_id or 'global'}-regression"
        )
        if runner.get_job(job_name) is None:
            runner.register_regression_job(
                job_name,
                agent=agent,
                workspace_id=workspace_id,
                agent_name=prompt_name,
                interval_seconds=policy.continuous_eval_interval_seconds,
                limit=policy.continuous_eval_limit,
                metadata={
                    "created_via": "self_learning_loop",
                    "self_learning_policy": {
                        **rollback_policy,
                        "experiment_id": experiment_id,
                        "target_version": target_version,
                    },
                    "task_type": task_type,
                    "user_segment": user_segment,
                },
            )
        return job_name

    def _resolve_reward_scorer(
        self,
        *,
        mode: str,
        workspace_id: str | None,
        prompt_name: str,
        limit: int,
    ) -> RewardScorer:
        return build_reward_scorer(
            mode,
            store=self.store,
            workspace_id=workspace_id,
            agent_name=prompt_name,
            limit=limit,
            fallback_scorer=self.reward_scorer,
        )

    def _historical_reward_score(self, examples: list[Any], scorer: RewardScorer) -> float:
        if not examples:
            return 0.0
        scores = [scorer.score_example(item).score for item in examples]
        return round(sum(scores) / len(scores), 4)

    def _serialize_generation_report(self, report: PromptGenerationReport) -> dict[str, Any]:
        return {
            "prompt_name": report.prompt_name,
            "workspace_id": report.workspace_id,
            "task_type": report.task_type,
            "user_segment": report.user_segment,
            "environment": report.environment,
            "based_on_version": report.based_on_version,
            "pattern_count": report.pattern_count,
            "lesson_count": report.lesson_count,
            "corrected_output_count": report.corrected_output_count,
            "candidates": [
                {
                    "content": candidate.content,
                    "strategy": candidate.strategy,
                    "rationale": candidate.rationale,
                    "metadata": dict(candidate.metadata),
                }
                for candidate in report.candidates
            ],
        }


__all__ = [
    "SelfLearningLoop",
    "SelfLearningLoopAction",
    "SelfLearningLoopPolicy",
    "SelfLearningLoopRun",
]
