"""Prompt candidate generation for self-learning loops."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .dataset import DatasetBuilder
from .improvement import FailurePattern, ReflectionLesson, ReflectionMemory
from .prompt_registry import PromptRecommendation, PromptRegistry
from .store import LearningStore
from uuid import uuid4


@dataclass
class PromptCandidate:
    """Generated prompt candidate with provenance."""

    content: str
    strategy: str
    rationale: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class PromptGenerationReport:
    """Structured report for one prompt evolution pass."""

    prompt_name: str
    workspace_id: str | None
    task_type: str | None
    user_segment: str | None
    environment: str
    based_on_version: int | None
    pattern_count: int
    lesson_count: int
    corrected_output_count: int
    candidates: list[PromptCandidate] = field(default_factory=list)


class PromptCandidateGenerator:
    """Build prompt candidates from production failures and lessons learned."""

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

    def generate(
        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,
        lessons: list[ReflectionLesson] | None = None,
        max_candidates: int = 5,
    ) -> list[PromptCandidate]:
        return self.generate_report(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            environment=environment,
            patterns=patterns,
            lessons=lessons,
            max_candidates=max_candidates,
        ).candidates

    def generate_report(
        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,
        lessons: list[ReflectionLesson] | None = None,
        max_candidates: int = 5,
    ) -> PromptGenerationReport:
        resolution = self.registry.resolve(
            prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            environment=environment,
        )
        versions = self.registry.list_versions(prompt_name)
        current_content = resolution.content if resolution is not None else (versions[0].content if versions else "")
        based_on_version = resolution.version if resolution is not None else (versions[0].version if versions else None)
        active_patterns = list(patterns or [])
        active_lessons = list(lessons or self._load_lessons(workspace_id=workspace_id, agent_name=prompt_name))
        corrected_outputs = self._corrected_outputs(workspace_id=workspace_id, agent_name=prompt_name)
        candidates: list[PromptCandidate] = []

        patch_lines = [pattern.suggested_fix for pattern in active_patterns[:3] if pattern.suggested_fix]
        if patch_lines:
            candidates.append(
                PromptCandidate(
                    content=self._append_lines(current_content, patch_lines),
                    strategy="patch",
                    rationale="Incorporate fixes from recurring failure clusters.",
                    metadata={"pattern_signatures": [pattern.signature for pattern in active_patterns[:3]]},
                )
            )

        if corrected_outputs:
            exemplar = corrected_outputs[0]
            candidates.append(
                PromptCandidate(
                    content=self._append_lines(
                        current_content,
                        [f"Use successful human edits as style exemplars and prefer outputs aligned with: {exemplar}"],
                    ),
                    strategy="style_correction",
                    rationale="Promote patterns observed in corrected human outputs.",
                    metadata={"corrected_output": exemplar},
                )
            )

        grounding_failures = {"hallucination", "bad_retrieval", "policy_violation"}
        if any(any(item in grounding_failures for item in pattern.failure_types) for pattern in active_patterns):
            candidates.append(
                PromptCandidate(
                    content=self._append_lines(
                        current_content,
                        [
                            "Answer only with claims grounded in available context or explicitly state uncertainty.",
                            "Prefer retrieved evidence and do not invent missing facts.",
                        ],
                    ),
                    strategy="retrieval_grounding",
                    rationale="Reduce hallucination and retrieval-grounding defects.",
                    metadata={"target_failures": sorted(grounding_failures)},
                )
            )

        efficiency_failures = {"latency", "too_expensive"}
        if any(any(item in efficiency_failures for item in pattern.failure_types) for pattern in active_patterns):
            candidates.append(
                PromptCandidate(
                    content=self._append_lines(
                        current_content,
                        [
                            "Keep the answer concise and avoid unnecessary reasoning or tool use.",
                            "Prefer the shortest correct answer that still satisfies the request.",
                        ],
                    ),
                    strategy="compress",
                    rationale="Reduce latency and token/cost overhead for this task class.",
                    metadata={"target_failures": sorted(efficiency_failures)},
                )
            )

        if active_lessons:
            lesson_lines = [lesson.suggested_action for lesson in active_lessons[:3] if lesson.suggested_action]
            if lesson_lines:
                candidates.append(
                    PromptCandidate(
                        content=self._rewrite_prompt(current_content, lesson_lines),
                        strategy="rewrite",
                        rationale="Consolidate reflection lessons into a cleaner operating prompt.",
                        metadata={"lesson_ids": [lesson.id for lesson in active_lessons[:3]]},
                    )
                )

        deduped: list[PromptCandidate] = []
        seen_contents = {current_content.strip()}
        for candidate in candidates:
            normalized = candidate.content.strip()
            if not normalized or normalized in seen_contents:
                continue
            seen_contents.add(normalized)
            deduped.append(candidate)
            if len(deduped) >= max(1, max_candidates):
                break
        return PromptGenerationReport(
            prompt_name=prompt_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            environment=environment,
            based_on_version=based_on_version,
            pattern_count=len(active_patterns),
            lesson_count=len(active_lessons),
            corrected_output_count=len(corrected_outputs),
            candidates=deduped,
        )

    def persist_candidates(self, report: PromptGenerationReport) -> list[PromptRecommendation]:
        records: list[PromptRecommendation] = []
        for index, candidate in enumerate(report.candidates, start=1):
            record = PromptRecommendation(
                id=uuid4().hex,
                prompt_name=report.prompt_name,
                kind="generated_candidate",
                title=f"Generated candidate {index} for {report.prompt_name}",
                summary=candidate.rationale,
                proposed_patch=candidate.content,
                based_on_version=report.based_on_version,
                metadata={
                    "strategy": candidate.strategy,
                    "workspace_id": report.workspace_id,
                    "task_type": report.task_type,
                    "user_segment": report.user_segment,
                    "environment": report.environment,
                    "candidate_metadata": dict(candidate.metadata),
                    "pattern_count": report.pattern_count,
                    "lesson_count": report.lesson_count,
                    "corrected_output_count": report.corrected_output_count,
                },
            )
            self.registry.save_recommendation(record)
            records.append(record)
        return records

    def _corrected_outputs(self, *, workspace_id: str | None, agent_name: str | None) -> list[str]:
        if self.store is None:
            return []
        builder = DatasetBuilder(self.store)
        examples = builder.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=100,
            include_unreviewed=False,
        )
        outputs: list[str] = []
        for example in examples:
            if example.corrected_output and example.corrected_output not in outputs:
                outputs.append(example.corrected_output)
        return outputs

    def _load_lessons(self, *, workspace_id: str | None, agent_name: str | None) -> list[ReflectionLesson]:
        if self.reflection_memory is None:
            return []
        return self.reflection_memory.list_lessons(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=20,
        )

    def _append_lines(self, current_content: str, lines: list[str]) -> str:
        base = current_content.rstrip()
        additions = "\n".join(f"- {line}" for line in lines if line)
        return f"{base}\n{additions}".strip()

    def _rewrite_prompt(self, current_content: str, lessons: list[str]) -> str:
        rules = "\n".join(f"{index + 1}. {lesson}" for index, lesson in enumerate(lessons) if lesson)
        if not rules:
            return current_content
        return (
            "Follow these operating rules.\n"
            f"{current_content.strip()}\n"
            "Refined directives:\n"
            f"{rules}"
        ).strip()


__all__ = ["PromptCandidate", "PromptCandidateGenerator", "PromptGenerationReport"]
