"""Policy learning mechanisms."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from logos_adk.learning.store import LearningStore

from .models import PolicyDecision
from .utils import _mean


@dataclass
class TrainingExport:
    """Exported training data for offline analysis."""

    workspace_id: str | None
    agent_name: str | None
    examples: list[dict[str, Any]]
    generated_at: str = field(default_factory=lambda: __import__('datetime').datetime.now(__import__('datetime').UTC).isoformat())
    metadata: dict[str, Any] = field(default_factory=dict)


class PolicyLearner:
    """Decision layer trained from existing runtime learning engines."""

    def __init__(
        self,
        store: "LearningStore",
        *,
        reward_model: Any | None = None,
        model_routing: Any | None = None,
        retrieval_learning: Any | None = None,
        tool_workflow: Any | None = None,
    ) -> None:
        """Initialize policy learner.
        
        Args:
            store: LearningStore instance
            reward_model: Optional reward model for scoring
            model_routing: Optional model routing engine
            retrieval_learning: Optional retrieval learning engine
            tool_workflow: Optional tool workflow learning engine
        """
        self.store = store
        self.reward_model = reward_model
        self.model_routing = model_routing
        self.retrieval_learning = retrieval_learning
        self.tool_workflow = tool_workflow

    def learn_policy(
        self,
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        limit: int = 500,
    ) -> PolicyDecision:
        """Learn policy decision from historical data."""
        runs = self.store.list_runs(
            workspace_id=workspace_id,
            agent_name=None,
            limit=limit,
        )
        
        if not runs:
            return self._default_policy(workspace_id, task_type, user_segment)
        
        reward_values = []
        model_counts: dict[str, int] = {}
        tool_counts: dict[str, int] = {}
        
        for run in runs:
            if self.reward_model:
                feedback = self.store.list_feedback(run_id=run.id)
                outcomes = self.store.list_outcomes(run_id=run.id)
                score = self.reward_model.score_run(run, feedback=feedback, outcomes=outcomes)
                reward_values.append(score.score)
            
            if run.model:
                model_counts[run.model] = model_counts.get(run.model, 0) + 1
            
            if run.tool_calls:
                for tool_call in run.tool_calls:
                    tool_counts[tool_call.get("name", "unknown")] = tool_counts.get(tool_call.get("name", "unknown"), 0) + 1
        
        avg_reward = _mean(reward_values, default=0.5)
        
        best_model = max(model_counts, key=model_counts.get) if model_counts else None
        preferred_tools = sorted(tool_counts, key=tool_counts.get, reverse=True)[:5] if tool_counts else []
        
        return PolicyDecision(
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            model_key=best_model,
            preferred_tools=preferred_tools,
            confidence=min(1.0, len(runs) / 100.0),
            average_reward=round(avg_reward, 4),
            rationale=self._build_rationale(runs, reward_values, model_counts, tool_counts),
            metrics={
                "sample_size": len(runs),
                "model_counts": model_counts,
                "tool_counts": tool_counts,
            },
        )

    def _default_policy(
        self,
        workspace_id: str | None,
        task_type: str | None,
        user_segment: str | None,
    ) -> PolicyDecision:
        """Return default policy when insufficient data."""
        return PolicyDecision(
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            confidence=0.0,
            average_reward=0.5,
            rationale=["Insufficient data for policy learning"],
            metrics={"sample_size": 0},
        )

    def _build_rationale(
        self,
        runs: list[Any],
        reward_values: list[float],
        model_counts: dict[str, int],
        tool_counts: dict[str, int],
    ) -> list[str]:
        """Build policy rationale from data."""
        rationale = []
        
        if reward_values:
            avg_reward = _mean(reward_values, default=0.5)
            if avg_reward >= 0.7:
                rationale.append(f"High average reward ({avg_reward:.2f}) indicates good policy")
            elif avg_reward <= 0.3:
                rationale.append(f"Low average reward ({avg_reward:.2f}) suggests policy issues")
        
        if model_counts:
            best_model = max(model_counts, key=model_counts.get)
            rationale.append(f"Model {best_model} most frequently used ({model_counts[best_model]} runs)")
        
        if tool_counts:
            best_tool = max(tool_counts, key=tool_counts.get)
            rationale.append(f"Tool {best_tool} most successful ({tool_counts[best_tool]} uses)")
        
        return rationale

    def export_training_data(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 500,
    ) -> TrainingExport:
        """Export training data for offline analysis."""
        from logos_adk.learning.dataset import DatasetBuilder
        
        builder = DatasetBuilder(self.store)
        examples = builder.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            include_unreviewed=False,
        )
        
        export_examples = []
        for example in examples:
            if self.reward_model:
                score = self.reward_model.score_example(example)
                export_examples.append({
                    "run_id": example.run_id,
                    "input_text": example.input_text,
                    "output_text": example.output_text,
                    "reward_score": score.score,
                    "reward_components": [
                        {"name": c.name, "value": c.raw_value, "weight": c.weight}
                        for c in score.components
                    ],
                    "metadata": example.metadata,
                })
        
        return TrainingExport(
            workspace_id=workspace_id,
            agent_name=agent_name,
            examples=export_examples,
            metadata={
                "sample_size": len(export_examples),
                "export_type": "policy_training",
            },
        )


__all__ = ["PolicyLearner", "TrainingExport"]
