"""Reward scoring mechanisms for learning."""
from __future__ import annotations

from dataclasses import asdict
from json import loads
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

if TYPE_CHECKING:
    from logos_adk.learning.models import LearningExample, OutcomeRecord, RunRecord, FeedbackRecord

from .models import RewardComponent, RewardScore, LearnedPreferenceModelState
from .utils import (
    _clip_unit,
    _mean,
    _preference_target,
    _reward_feature_values,
    _sigmoid,
)


@runtime_checkable
class RewardScorer(Protocol):
    """Shared interface for heuristic and learned reward scorers."""

    def score_example(self, example: "LearningExample") -> RewardScore:
        """Return a unified reward score for a single example."""
        ...

    def score_run(
        self,
        run: "RunRecord",
        *,
        feedback: list["FeedbackRecord"],
        outcomes: list["OutcomeRecord"],
    ) -> RewardScore:
        """Return a unified reward score for a captured run."""
        ...


class HeuristicRewardScorer:
    """Aggregate human feedback, outcomes, cost, latency, and safety into one score."""

    def __init__(
        self,
        *,
        human_weight: float = 0.45,
        outcome_weight: float = 0.35,
        cost_weight: float = 0.08,
        latency_weight: float = 0.06,
        safety_weight: float = 0.06,
    ) -> None:
        self.human_weight = human_weight
        self.outcome_weight = outcome_weight
        self.cost_weight = cost_weight
        self.latency_weight = latency_weight
        self.safety_weight = safety_weight

    def score_example(self, example: "LearningExample") -> RewardScore:
        """Score a learning example using heuristic weights."""
        features = _reward_feature_values(example)
        components = [
            RewardComponent("human_feedback", features["human_feedback"], self.human_weight, features["human_feedback"] * self.human_weight),
            RewardComponent("business_outcome", features["business_outcome"], self.outcome_weight, features["business_outcome"] * self.outcome_weight),
            RewardComponent("cost_penalty", features["cost_penalty"], -self.cost_weight, -features["cost_penalty"] * self.cost_weight),
            RewardComponent("latency_penalty", features["latency_penalty"], -self.latency_weight, -features["latency_penalty"] * self.latency_weight),
            RewardComponent("safety_penalty", features["safety_penalty"], -self.safety_weight, -features["safety_penalty"] * self.safety_weight),
        ]
        score = max(0.0, min(1.0, sum(component.contribution for component in components)))
        return RewardScore(
            score=round(score, 4),
            trace_id=example.trace_id,
            run_id=example.run_id,
            workspace_id=example.workspace_id,
            agent_name=example.agent_name,
            components=components,
        )

    def score_run(
        self,
        run: "RunRecord",
        *,
        feedback: list["FeedbackRecord"],
        outcomes: list["OutcomeRecord"],
    ) -> RewardScore:
        """Score a run using heuristic weights."""
        from logos_adk.learning.models import LearningExample
        
        example = LearningExample(
            run_id=run.id,
            trace_id=run.trace_id,
            agent_name=run.agent_name,
            input_text=run.input_text,
            output_text=run.output_text,
            created_at=run.created_at,
            session_id=run.session_id,
            workspace_id=run.workspace_id,
            prompt_version=run.prompt_version,
            model=run.model,
            config_fingerprint=run.config_fingerprint,
            feedback=feedback,
            outcomes=outcomes,
            failure_types=sorted({failure for item in feedback for failure in item.failure_types}),
            corrected_output=next((item.corrected_output for item in feedback if item.corrected_output), None),
            metadata={**run.metadata, "usage": run.usage or {}},
        )
        return self.score_example(example)


class RewardModel(HeuristicRewardScorer):
    """Backward-compatible alias for the heuristic reward model."""


class LearnedPreferenceScorer:
    """Lightweight learned preference model trained from observed preference signals."""

    MODEL_VERSION = 1
    DEFAULT_FEATURE_WEIGHTS: dict[str, float] = {
        "human_feedback": 1.0,
        "business_outcome": 0.8,
        "cost_penalty": -0.25,
        "latency_penalty": -0.2,
        "safety_penalty": -1.0,
        "correction_penalty": -0.7,
    }

    def __init__(
        self,
        *,
        bias: float = 0.0,
        feature_weights: dict[str, float] | None = None,
        training_summary: dict[str, Any] | None = None,
        fallback_scorer: RewardScorer | None = None,
    ) -> None:
        self.bias = float(bias)
        self.feature_weights = dict(feature_weights or self.DEFAULT_FEATURE_WEIGHTS)
        self.training_summary = dict(training_summary or {})
        self.fallback_scorer = fallback_scorer or HeuristicRewardScorer()

    def score_example(self, example: "LearningExample") -> RewardScore:
        """Score example using learned preference model."""
        features = _reward_feature_values(example)
        contributions: list[RewardComponent] = []
        linear_score = self.bias
        for name, weight in self.feature_weights.items():
            raw_value = features.get(name, 0.0)
            centered = raw_value - 0.5
            contribution = centered * float(weight)
            linear_score += contribution
            contributions.append(
                RewardComponent(
                    name=name,
                    raw_value=raw_value,
                    weight=float(weight),
                    contribution=round(contribution, 6),
                    metadata={"centered_value": round(centered, 6), "scorer": "learned_preference"},
                )
            )
        score = _clip_unit(_sigmoid(linear_score))
        return RewardScore(
            score=round(score, 4),
            trace_id=example.trace_id,
            run_id=example.run_id,
            workspace_id=example.workspace_id,
            agent_name=example.agent_name,
            components=contributions,
        )

    def score_run(
        self,
        run: "RunRecord",
        *,
        feedback: list["FeedbackRecord"],
        outcomes: list["OutcomeRecord"],
    ) -> RewardScore:
        """Score run using learned preference model."""
        from logos_adk.learning.models import LearningExample
        
        example = LearningExample(
            run_id=run.id,
            trace_id=run.trace_id,
            agent_name=run.agent_name,
            input_text=run.input_text,
            output_text=run.output_text,
            created_at=run.created_at,
            session_id=run.session_id,
            workspace_id=run.workspace_id,
            prompt_version=run.prompt_version,
            model=run.model,
            config_fingerprint=run.config_fingerprint,
            feedback=feedback,
            outcomes=outcomes,
            failure_types=sorted({failure for item in feedback for failure in item.failure_types}),
            corrected_output=next((item.corrected_output for item in feedback if item.corrected_output), None),
            metadata={**run.metadata, "usage": run.usage or {}},
        )
        return self.score_example(example)

    @classmethod
    def train_from_examples(
        cls,
        examples: list["LearningExample"],
        *,
        fallback_scorer: RewardScorer | None = None,
    ) -> "LearnedPreferenceScorer":
        """Train preference scorer from examples."""
        if not examples:
            return cls(fallback_scorer=fallback_scorer, training_summary={"sample_size": 0, "mode": "fallback"})

        features_by_name = sorted(cls.DEFAULT_FEATURE_WEIGHTS)
        positives = {name: [] for name in features_by_name}
        negatives = {name: [] for name in features_by_name}
        labels: list[float] = []
        preference_pairs = 0

        for example in examples:
            label = _preference_target(example)
            labels.append(label)
            features = _reward_feature_values(example)
            bucket = positives if label >= 0.5 else negatives
            certainty = max(0.05, abs(label - 0.5) * 2.0)
            for name in features_by_name:
                bucket[name].append(features.get(name, 0.0) * certainty)
            if example.corrected_output or any(item.expected_output for item in example.feedback):
                preference_pairs += 1

        feature_weights: dict[str, float] = {}
        for name in features_by_name:
            pos_mean = _mean(positives[name], default=0.0)
            neg_mean = _mean(negatives[name], default=0.0)
            weight = pos_mean - neg_mean
            if abs(weight) < 0.05:
                weight = cls.DEFAULT_FEATURE_WEIGHTS.get(name, 0.0)
            feature_weights[name] = round(weight, 4)

        mean_label = _clip_unit(_mean(labels, default=0.5))
        clipped_mean = min(0.99, max(0.01, mean_label))
        from math import log
        bias = log(clipped_mean / (1.0 - clipped_mean))
        return cls(
            bias=round(bias, 4),
            feature_weights=feature_weights,
            fallback_scorer=fallback_scorer,
            training_summary={
                "sample_size": len(examples),
                "positive_examples": sum(1 for item in labels if item >= 0.5),
                "negative_examples": sum(1 for item in labels if item < 0.5),
                "mean_label": round(mean_label, 4),
                "preference_pairs": preference_pairs,
            },
        )

    @classmethod
    def train_from_store(
        cls,
        store: Any,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 500,
        fallback_scorer: RewardScorer | None = None,
    ) -> "LearnedPreferenceScorer":
        """Train preference scorer from learning store."""
        from logos_adk.learning.dataset import DatasetBuilder
        
        builder = DatasetBuilder(store)
        examples = builder.build(
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            include_unreviewed=False,
        )
        return cls.train_from_examples(examples, fallback_scorer=fallback_scorer)

    def export_state(self) -> LearnedPreferenceModelState:
        """Export scorer state."""
        return LearnedPreferenceModelState(
            version=self.MODEL_VERSION,
            bias=round(self.bias, 6),
            feature_weights={key: round(float(value), 6) for key, value in self.feature_weights.items()},
            training_summary=dict(self.training_summary),
        )

    @classmethod
    def from_state(
        cls,
        state: LearnedPreferenceModelState | dict[str, Any],
        *,
        fallback_scorer: RewardScorer | None = None,
    ) -> "LearnedPreferenceScorer":
        """Create scorer from exported state."""
        payload = asdict(state) if isinstance(state, LearnedPreferenceModelState) else dict(state)
        return cls(
            bias=float(payload.get("bias", 0.0) or 0.0),
            feature_weights=dict(payload.get("feature_weights") or {}),
            training_summary=dict(payload.get("training_summary") or {}),
            fallback_scorer=fallback_scorer,
        )

    def save(self, path: str) -> str:
        """Save scorer state to file."""
        from json import dumps
        from pathlib import Path
        
        destination = Path(path)
        if destination.parent and not destination.parent.exists():
            destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_text(
            dumps(asdict(self.export_state()), ensure_ascii=False, indent=2),
            encoding="utf-8"
        )
        return str(destination)

    @classmethod
    def load(
        cls,
        path: str,
        *,
        fallback_scorer: RewardScorer | None = None,
    ) -> "LearnedPreferenceScorer":
        """Load scorer state from file."""
        payload = loads(Path(path).read_text(encoding="utf-8"))
        return cls.from_state(payload, fallback_scorer=fallback_scorer)


def build_reward_scorer(
    mode: str,
    *,
    store: Any | None = None,
    workspace_id: str | None = None,
    agent_name: str | None = None,
    limit: int = 500,
    fallback_scorer: RewardScorer | None = None,
) -> RewardScorer:
    """Construct a reward scorer for the requested mode."""
    normalized = str(mode or "heuristic").strip().lower()
    if normalized == "heuristic":
        return RewardModel()
    if normalized == "learned":
        if store is None:
            raise ValueError("store is required for learned reward scorer mode")
        return LearnedPreferenceScorer.train_from_store(
            store,
            workspace_id=workspace_id,
            agent_name=agent_name,
            limit=limit,
            fallback_scorer=fallback_scorer,
        )
    raise ValueError(f"Unsupported reward scorer mode: {mode}")


__all__ = [
    "RewardScorer",
    "HeuristicRewardScorer",
    "RewardModel",
    "LearnedPreferenceScorer",
    "build_reward_scorer",
]
