"""Build learning datasets from captured production runs."""
from __future__ import annotations

from dataclasses import asdict
from typing import Any

from .models import LearningExample
from .store import LearningStore


class DatasetBuilder:
    """Join runs, feedback, and outcomes into learning examples."""

    def __init__(self, store: LearningStore) -> None:
        self.store = store

    def build(
        self,
        *,
        workspace_id: str | None = None,
        agent_name: str | None = None,
        limit: int = 100,
        include_unreviewed: bool = True,
    ) -> list[LearningExample]:
        runs = self.store.list_runs(workspace_id=workspace_id, agent_name=agent_name, limit=limit)
        examples: list[LearningExample] = []
        for run in runs:
            feedback = self.store.list_feedback(run_id=run.id, trace_id=run.trace_id, workspace_id=workspace_id)
            outcomes = self.store.list_outcomes(run_id=run.id, trace_id=run.trace_id, workspace_id=workspace_id)
            if not include_unreviewed and not feedback and not outcomes:
                continue
            failure_types: list[str] = []
            corrected_output = None
            for item in feedback:
                for failure_type in item.failure_types:
                    if failure_type not in failure_types:
                        failure_types.append(failure_type)
                if corrected_output is None and item.corrected_output:
                    corrected_output = item.corrected_output
            examples.append(
                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=failure_types,
                    corrected_output=corrected_output,
                    metadata={
                        **dict(run.metadata),
                        "usage": dict(run.usage or {}),
                    },
                )
            )
        return examples

    def build_payload(self, **kwargs: Any) -> dict[str, Any]:
        """Return a JSON-friendly dataset payload."""
        items = self.build(**kwargs)
        return {
            "summary": {
                "count": len(items),
                "with_feedback": sum(1 for item in items if item.feedback),
                "with_outcomes": sum(1 for item in items if item.outcomes),
            },
            "items": [asdict(item) for item in items],
        }
