"""Cross-agent learning and safe transfer of reusable artifacts."""
from __future__ import annotations

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

from .dataset import DatasetBuilder
from .improvement import ReflectionMemory, RewardModel, RewardScorer
from .models import utcnow
from .prompt_registry import PromptRecommendation, PromptRegistry
from .retrieval_learning import RetrievalLearningEngine
from .store import LearningStore
from .tool_workflow_learning import ToolWorkflowLearningEngine


@dataclass
class TransferArtifact:
    """Reusable learning artifact that may transfer across agents."""

    id: str
    kind: str
    source_agent_name: str
    workspace_id: str | None
    task_type: str | None
    user_segment: str | None
    scope: str
    title: str
    summary: str
    affected_surface: str
    confidence: float
    payload: dict[str, Any] = field(default_factory=dict)
    guardrails: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class TransferRecommendation:
    """Recommended transfer from one agent into another."""

    artifact_id: str
    kind: str
    source_agent_name: str
    target_agent_name: str
    workspace_id: str | None
    compatibility_score: float
    estimated_uplift: float
    status: str
    reason: str
    affected_surface: str
    safe_apply: bool
    payload: dict[str, Any] = field(default_factory=dict)


@dataclass
class TransferApplyResult:
    """Result of applying or previewing cross-agent transfer recommendations."""

    source_agent_name: str
    target_agent_name: str
    workspace_id: str | None
    safe_apply: bool
    recommended: int
    applied: int
    created_lessons: int = 0
    created_prompt_recommendations: int = 0
    items: list[dict[str, Any]] = field(default_factory=list)


@dataclass
class TransferUpliftReport:
    """Separate uplift measurement for transferred vs local learning."""

    workspace_id: str | None
    target_agent_name: str
    generated_at: str
    transferred_runs: int
    local_runs: int
    transferred_average_reward: float
    local_average_reward: float
    delta: float


def _guardrails(value: Any) -> list[str]:
    if isinstance(value, list):
        return [str(item) for item in value]
    return []


class KnowledgeTransferEngine:
    """Collect, recommend, and safely apply cross-agent learning artifacts."""

    def __init__(
        self,
        store: LearningStore,
        *,
        reflection_memory: ReflectionMemory,
        prompt_registry: PromptRegistry | None = None,
        retrieval_learning: RetrievalLearningEngine | None = None,
        tool_workflow: ToolWorkflowLearningEngine | None = None,
        model_routing: Any | None = None,
        reward_scorer: RewardScorer | None = None,
    ) -> None:
        self.store = store
        self.datasets = DatasetBuilder(store)
        self.reflection_memory = reflection_memory
        self.prompt_registry = prompt_registry
        self.retrieval_learning = retrieval_learning
        self.tool_workflow = tool_workflow
        self.model_routing = model_routing
        self.reward_scorer = reward_scorer or RewardModel()

    def collect_artifacts(
        self,
        *,
        source_agent_name: str,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        limit: int = 50,
    ) -> list[TransferArtifact]:
        items: list[TransferArtifact] = []

        for lesson in self.reflection_memory.list_lessons(
            workspace_id=workspace_id,
            agent_name=source_agent_name,
            limit=limit,
        ):
            items.append(
                TransferArtifact(
                    id=f"lesson:{lesson.id}",
                    kind="reflection_lesson",
                    source_agent_name=source_agent_name,
                    workspace_id=lesson.workspace_id,
                    task_type=lesson.metadata.get("task_type"),
                    user_segment=lesson.metadata.get("user_segment"),
                    scope=str(lesson.metadata.get("scope", "agent_local")),
                    title=lesson.lesson,
                    summary=lesson.rationale,
                    affected_surface=str(lesson.category),
                    confidence=float(lesson.confidence),
                    payload=asdict(lesson),
                    guardrails=_guardrails(lesson.metadata.get("guardrails")),
                    metadata=dict(lesson.metadata),
                )
            )

        if self.prompt_registry is not None:
            for recommendation in self.prompt_registry.list_recommendations(source_agent_name):
                if recommendation.kind not in {"prompt_patch", "generated_candidate"}:
                    continue
                items.append(
                    TransferArtifact(
                        id=f"prompt:{recommendation.id}",
                        kind="prompt_patch",
                        source_agent_name=source_agent_name,
                        workspace_id=recommendation.metadata.get("workspace_id", workspace_id),
                        task_type=recommendation.metadata.get("task_type"),
                        user_segment=recommendation.metadata.get("user_segment"),
                        scope=str(recommendation.metadata.get("scope", "agent_local")),
                        title=recommendation.title,
                        summary=recommendation.summary,
                        affected_surface="prompt",
                        confidence=min(0.95, 0.4 + sum(recommendation.failure_counts.values()) * 0.05),
                        payload=asdict(recommendation),
                        guardrails=_guardrails(recommendation.metadata.get("guardrails")),
                        metadata=dict(recommendation.metadata),
                    )
                )

        if self.retrieval_learning is not None:
            profiles = self.retrieval_learning.list_profiles(
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
            )
            for profile in profiles[:5]:
                items.append(
                    TransferArtifact(
                        id=f"retrieval:{profile.workspace_id}:{profile.task_type}:{profile.user_segment}:{profile.limit}",
                        kind="retrieval_profile",
                        source_agent_name=source_agent_name,
                        workspace_id=profile.workspace_id,
                        task_type=profile.task_type,
                        user_segment=profile.user_segment,
                        scope="shared",
                        title="Shared retrieval profile",
                        summary="High-performing retrieval settings for this segment.",
                        affected_surface="retrieval_policy",
                        confidence=float(profile.score),
                        payload=asdict(profile),
                        metadata={"scope": "shared"},
                    )
                )

        if self.tool_workflow is not None:
            tool_scores = self.tool_workflow.list_tool_scores(workspace_id=workspace_id, task_type=task_type)
            if tool_scores:
                items.append(
                    TransferArtifact(
                        id=f"tools:{workspace_id}:{task_type or '*'}",
                        kind="tool_preference",
                        source_agent_name=source_agent_name,
                        workspace_id=workspace_id,
                        task_type=task_type,
                        user_segment=user_segment,
                        scope="shared",
                        title="Shared tool preference set",
                        summary="Historically successful tools for this task segment.",
                        affected_surface="tool_policy",
                        confidence=max(float(item.score) for item in tool_scores[:3]),
                        payload={"tools": [asdict(item) for item in tool_scores[:5]]},
                        metadata={"scope": "shared"},
                    )
                )

            fallback_scores = self.tool_workflow.list_fallback_scores(workspace_id=workspace_id, task_type=task_type)
            if fallback_scores:
                items.append(
                    TransferArtifact(
                        id=f"routing:{workspace_id}:{task_type or '*'}",
                        kind="routing_policy",
                        source_agent_name=source_agent_name,
                        workspace_id=workspace_id,
                        task_type=task_type,
                        user_segment=user_segment,
                        scope="shared",
                        title="Shared routing fallback policy",
                        summary="Historically successful fallback ordering for this segment.",
                        affected_surface="routing_policy",
                        confidence=max(float(item.score) for item in fallback_scores[:3]),
                        payload={"fallbacks": [asdict(item) for item in fallback_scores[:5]]},
                        metadata={"scope": "shared"},
                    )
                )

        if self.model_routing is not None:
            model_scores = self.model_routing.list_model_scores(workspace_id=workspace_id, task_type=task_type)
            if model_scores:
                items.append(
                    TransferArtifact(
                        id=f"model-routing:{workspace_id}:{task_type or '*'}",
                        kind="routing_policy",
                        source_agent_name=source_agent_name,
                        workspace_id=workspace_id,
                        task_type=task_type,
                        user_segment=user_segment,
                        scope="shared",
                        title="Shared model routing policy",
                        summary="Historically successful model routing for this task segment.",
                        affected_surface="routing_policy",
                        confidence=max(float(item.score) for item in model_scores[:3]),
                        payload={"models": [asdict(item) for item in model_scores[:5]]},
                        metadata={"scope": "shared"},
                    )
                )

        items.sort(key=lambda item: (item.confidence, item.kind), reverse=True)
        return items[:limit]

    def recommend_transfers(
        self,
        *,
        source_agent_name: str,
        target_agent_name: str,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        target_guardrails: list[str] | None = None,
        safe_apply: bool = True,
        limit: int = 20,
    ) -> list[TransferRecommendation]:
        artifacts = self.collect_artifacts(
            source_agent_name=source_agent_name,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            limit=max(limit * 2, 20),
        )
        recommendations: list[TransferRecommendation] = []
        for artifact in artifacts:
            compatibility, reason = self._compatibility(
                artifact,
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
                target_guardrails=target_guardrails or [],
            )
            if compatibility < 0.45:
                continue
            estimated_uplift = round(max(0.0, (compatibility - 0.4) * max(0.1, artifact.confidence)), 4)
            recommendations.append(
                TransferRecommendation(
                    artifact_id=artifact.id,
                    kind=artifact.kind,
                    source_agent_name=source_agent_name,
                    target_agent_name=target_agent_name,
                    workspace_id=workspace_id,
                    compatibility_score=compatibility,
                    estimated_uplift=estimated_uplift,
                    status="recommended" if safe_apply else "ready_to_apply",
                    reason=reason,
                    affected_surface=artifact.affected_surface,
                    safe_apply=safe_apply,
                    payload=asdict(artifact),
                )
            )
        recommendations.sort(key=lambda item: (item.compatibility_score, item.estimated_uplift), reverse=True)
        return recommendations[:limit]

    def apply_recommendations(
        self,
        recommendations: list[TransferRecommendation],
        *,
        safe_apply: bool = True,
    ) -> TransferApplyResult:
        if not recommendations:
            return TransferApplyResult(
                source_agent_name="",
                target_agent_name="",
                workspace_id=None,
                safe_apply=safe_apply,
                recommended=0,
                applied=0,
            )
        source_agent_name = recommendations[0].source_agent_name
        target_agent_name = recommendations[0].target_agent_name
        workspace_id = recommendations[0].workspace_id
        result = TransferApplyResult(
            source_agent_name=source_agent_name,
            target_agent_name=target_agent_name,
            workspace_id=workspace_id,
            safe_apply=safe_apply,
            recommended=len(recommendations),
            applied=0,
        )
        for item in recommendations:
            if safe_apply:
                result.items.append(
                    {
                        "artifact_id": item.artifact_id,
                        "status": "recommended_only",
                        "reason": item.reason,
                        "affected_surface": item.affected_surface,
                    }
                )
                continue
            artifact = item.payload
            kind = str(item.kind)
            if kind == "prompt_patch" and self.prompt_registry is not None:
                recommendation_payload = artifact.get("payload", {})
                self.prompt_registry.ensure_prompt(target_agent_name)
                record = PromptRecommendation(
                    id=uuid4().hex,
                    prompt_name=target_agent_name,
                    kind="transfer_patch",
                    title=f"Transferred patch from {source_agent_name}",
                    summary=str(artifact.get("summary") or "Transferred cross-agent prompt recommendation."),
                    proposed_patch=recommendation_payload.get("proposed_patch"),
                    based_on_version=recommendation_payload.get("based_on_version"),
                    failure_counts=dict(recommendation_payload.get("failure_counts") or {}),
                    metadata={
                        "transfer_source_agent": source_agent_name,
                        "transfer_compatibility_score": item.compatibility_score,
                        "workspace_id": workspace_id,
                    },
                )
                self.prompt_registry.save_recommendation(record)
                result.created_prompt_recommendations += 1
                result.applied += 1
                result.items.append({"artifact_id": item.artifact_id, "status": "applied", "kind": kind})
                continue

            lesson_text, action = self._transfer_lesson_payload(item)
            self.reflection_memory.create_lesson(
                workspace_id=workspace_id,
                agent_name=target_agent_name,
                category=item.affected_surface,
                lesson=lesson_text,
                rationale=f"Transferred from {source_agent_name}. {item.reason}",
                suggested_action=action,
                confidence=item.compatibility_score,
                metadata={
                    "transfer_source_agent": source_agent_name,
                    "transfer_artifact_kind": kind,
                    "transfer_safe_apply": False,
                    "scope": "shared",
                },
            )
            result.created_lessons += 1
            result.applied += 1
            result.items.append({"artifact_id": item.artifact_id, "status": "applied", "kind": kind})
        return result

    def measure_transfer_uplift(
        self,
        *,
        target_agent_name: str,
        workspace_id: str | None = None,
        limit: int = 200,
    ) -> TransferUpliftReport:
        examples = self.datasets.build(
            workspace_id=workspace_id,
            agent_name=target_agent_name,
            limit=limit,
            include_unreviewed=False,
        )
        transferred_scores: list[float] = []
        local_scores: list[float] = []
        for example in examples:
            metadata = example.metadata if isinstance(example.metadata, dict) else {}
            score = self.reward_scorer.score_example(example).score
            if metadata.get("transfer_source_agent"):
                transferred_scores.append(score)
            else:
                local_scores.append(score)
        transferred_average = round(sum(transferred_scores) / len(transferred_scores), 4) if transferred_scores else 0.0
        local_average = round(sum(local_scores) / len(local_scores), 4) if local_scores else 0.0
        return TransferUpliftReport(
            workspace_id=workspace_id,
            target_agent_name=target_agent_name,
            generated_at=utcnow(),
            transferred_runs=len(transferred_scores),
            local_runs=len(local_scores),
            transferred_average_reward=transferred_average,
            local_average_reward=local_average,
            delta=round(transferred_average - local_average, 4),
        )

    def _compatibility(
        self,
        artifact: TransferArtifact,
        *,
        workspace_id: str | None,
        task_type: str | None,
        user_segment: str | None,
        target_guardrails: list[str],
    ) -> tuple[float, str]:
        score = 0.0
        reasons: list[str] = []
        if artifact.scope == "shared":
            score += 0.2
            reasons.append("shared_scope")
        else:
            score += 0.1
            reasons.append("agent_local_scope")
        if artifact.workspace_id is None or artifact.workspace_id == workspace_id:
            score += 0.25
            reasons.append("workspace_compatible")
        if artifact.task_type is None or artifact.task_type == task_type:
            score += 0.2
            reasons.append("task_type_compatible")
        if artifact.user_segment is None or artifact.user_segment == user_segment:
            score += 0.1
            reasons.append("user_segment_compatible")
        if not artifact.guardrails or not target_guardrails:
            score += 0.15
            reasons.append("guardrails_unspecified")
        elif set(artifact.guardrails) & set(target_guardrails):
            score += 0.15
            reasons.append("guardrails_compatible")
        else:
            reasons.append("guardrails_mismatch")
            score -= 0.15
        score += min(0.25, max(0.0, float(artifact.confidence)) * 0.25)
        return round(max(0.0, min(1.0, score)), 4), ", ".join(reasons)

    def _transfer_lesson_payload(self, recommendation: TransferRecommendation) -> tuple[str, str]:
        artifact = recommendation.payload
        payload = artifact.get("payload", {})
        kind = recommendation.kind
        if kind == "reflection_lesson":
            return (
                f"Transferred lesson from {recommendation.source_agent_name}: {artifact.get('title')}",
                str(payload.get("suggested_action") or artifact.get("summary") or "Review and adapt the transferred lesson."),
            )
        if kind == "retrieval_profile":
            return (
                f"Transferred retrieval profile from {recommendation.source_agent_name}",
                "Adopt or experiment with the transferred retrieval settings for the matching segment.",
            )
        if kind == "tool_preference":
            return (
                f"Transferred tool preference set from {recommendation.source_agent_name}",
                "Prefer the transferred tool ordering for the matching task segment.",
            )
        if kind == "routing_policy":
            return (
                f"Transferred routing policy from {recommendation.source_agent_name}",
                "Adopt the transferred fallback/routing preferences after validation.",
            )
        return (
            f"Transferred artifact from {recommendation.source_agent_name}",
            "Review and apply the transferred recommendation.",
        )


__all__ = [
    "KnowledgeTransferEngine",
    "TransferApplyResult",
    "TransferArtifact",
    "TransferRecommendation",
    "TransferUpliftReport",
]
