"""
Training Data Exporter

Exports logged interactions and feedback into formats suitable for
model training (SFT, DPO, RLHF).
"""

import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any

from sage.data_collection.schemas import (
    ExpertFeedback,
    FeedbackRating,
    LoggedInteraction,
    TrainingExample,
)

logger = logging.getLogger(__name__)


class TrainingDataExporter:
    """Exports interaction data for model training.

    Supports multiple output formats:
    - SFT (Supervised Fine-Tuning): prompt -> response pairs
    - DPO (Direct Preference Optimization): prompt -> chosen/rejected pairs
    - JSONL: Raw interaction data

    Example:
        exporter = TrainingDataExporter()

        # Export for SFT
        sft_data = exporter.prepare_sft_data(interactions, feedback)
        exporter.export_sft("sft_training.jsonl", sft_data)

        # Export for DPO
        dpo_data = exporter.prepare_preference_data(interactions, feedback)
        exporter.export_dpo("dpo_training.jsonl", dpo_data)
    """

    def __init__(
        self,
        include_context: bool = True,
        min_confidence: float = 0.5,
        require_validation: bool = False,
    ):
        """Initialize exporter.

        Args:
            include_context: Include RAG context in prompts.
            min_confidence: Minimum confidence score for inclusion.
            require_validation: Require expert validation for inclusion.
        """
        self.include_context = include_context
        self.min_confidence = min_confidence
        self.require_validation = require_validation

    def prepare_sft_data(
        self,
        interactions: list[LoggedInteraction],
        feedback_map: dict[str, ExpertFeedback] | None = None,
        min_score: int = 4,
    ) -> list[TrainingExample]:
        """Prepare interactions for supervised fine-tuning.

        Args:
            interactions: List of logged interactions.
            feedback_map: Map of interaction_id -> feedback.
            min_score: Minimum feedback score for inclusion.

        Returns:
            List of TrainingExample objects.
        """
        examples = []
        feedback_map = feedback_map or {}

        for interaction in interactions:
            # Skip low confidence
            if interaction.response_confidence < self.min_confidence:
                continue

            # Skip abstained/error outcomes
            if interaction.outcome.value in ("abstained", "error"):
                continue

            # Check validation if required
            if self.require_validation:
                feedback = feedback_map.get(interaction.interaction_id)
                if not feedback or feedback.accuracy_score.value < min_score:
                    continue

            # Build context string
            context = None
            if self.include_context and interaction.rag_context:
                context_parts = []
                for ctx in interaction.rag_context[:3]:  # Limit context
                    context_parts.append(
                        f"[Source: {ctx.doc_title}"
                        f"{', ' + ctx.section if ctx.section else ''}]\n"
                        f"{ctx.content[:500]}"
                    )
                context = "\n\n".join(context_parts)

            # Get scores from feedback if available
            feedback = feedback_map.get(interaction.interaction_id)
            accuracy_score = feedback.accuracy_score.value / 5 if feedback else 0.0
            conservatism_score = feedback.conservatism_score.value / 5 if feedback else 0.0

            # Use corrected response if available
            response = interaction.response
            if feedback and feedback.corrected_response:
                response = feedback.corrected_response

            example = TrainingExample(
                prompt=interaction.query,
                response=response,
                context=context,
                expert_validated=feedback is not None,
                accuracy_score=accuracy_score,
                conservatism_score=conservatism_score,
                source_interaction_id=interaction.interaction_id,
                source_feedback_id=feedback.feedback_id if feedback else None,
                tags=self._extract_tags(interaction),
            )

            examples.append(example)

        logger.info(f"Prepared {len(examples)} SFT examples from {len(interactions)} interactions")
        return examples

    def prepare_preference_data(
        self,
        interactions: list[LoggedInteraction],
        feedback_map: dict[str, ExpertFeedback],
    ) -> list[TrainingExample]:
        """Prepare data for preference learning (DPO/RLHF).

        Creates pairs where:
        - chosen: corrected or high-rated response
        - rejected: original response (if corrected) or low-rated

        Args:
            interactions: List of logged interactions.
            feedback_map: Map of interaction_id -> feedback.

        Returns:
            List of TrainingExample objects with preference data.
        """
        examples = []

        # Group interactions by query similarity for contrast pairs
        interaction_map = {i.interaction_id: i for i in interactions}

        for interaction_id, feedback in feedback_map.items():
            interaction = interaction_map.get(interaction_id)
            if not interaction:
                continue

            # Skip if no correction available
            if not feedback.corrected_response:
                continue

            # Build context
            context = None
            if self.include_context and interaction.rag_context:
                context_parts = []
                for ctx in interaction.rag_context[:3]:
                    context_parts.append(
                        f"[Source: {ctx.doc_title}]\n{ctx.content[:500]}"
                    )
                context = "\n\n".join(context_parts)

            example = TrainingExample(
                prompt=interaction.query,
                response=feedback.corrected_response,  # Use corrected as main response
                context=context,
                preferred_response=feedback.corrected_response,
                rejected_response=interaction.response,  # Original is rejected
                expert_validated=True,
                accuracy_score=feedback.accuracy_score.value / 5,
                conservatism_score=feedback.conservatism_score.value / 5,
                source_interaction_id=interaction_id,
                source_feedback_id=feedback.feedback_id,
                tags=self._extract_tags(interaction) + ["preference_pair"],
            )

            examples.append(example)

        logger.info(f"Prepared {len(examples)} preference pairs")
        return examples

    def _extract_tags(self, interaction: LoggedInteraction) -> list[str]:
        """Extract tags from interaction for categorization."""
        tags = []

        # Query type
        tags.append(f"type:{interaction.query_type}")

        # Complexity
        tags.append(f"complexity:{interaction.complexity}")

        # Tool usage
        if interaction.tool_executions:
            tags.append("uses_tools")
            for tool in interaction.tool_executions:
                tags.append(f"tool:{tool.tool_name}")

        # Safety
        if interaction.abstention_triggered:
            tags.append("abstained")
        if interaction.escalation_triggered:
            tags.append("escalated")

        # Confidence bands
        if interaction.response_confidence >= 0.9:
            tags.append("high_confidence")
        elif interaction.response_confidence >= 0.7:
            tags.append("medium_confidence")
        else:
            tags.append("low_confidence")

        return tags

    def export_sft(
        self,
        output_path: str | Path,
        examples: list[TrainingExample],
        format: str = "jsonl",
    ) -> int:
        """Export SFT training data.

        Args:
            output_path: Path to output file.
            examples: List of TrainingExample objects.
            format: Output format (jsonl, json, alpaca).

        Returns:
            Number of examples exported.
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if format == "jsonl":
            with open(output_path, "w") as f:
                for example in examples:
                    f.write(json.dumps(example.to_sft_format()) + "\n")

        elif format == "json":
            with open(output_path, "w") as f:
                json.dump([e.to_sft_format() for e in examples], f, indent=2)

        elif format == "alpaca":
            # Alpaca format: instruction, input, output
            alpaca_data = []
            for example in examples:
                entry = {
                    "instruction": example.prompt,
                    "input": example.context or "",
                    "output": example.response,
                }
                alpaca_data.append(entry)

            with open(output_path, "w") as f:
                json.dump(alpaca_data, f, indent=2)

        else:
            raise ValueError(f"Unknown format: {format}")

        logger.info(f"Exported {len(examples)} SFT examples to {output_path}")
        return len(examples)

    def export_dpo(
        self,
        output_path: str | Path,
        examples: list[TrainingExample],
    ) -> int:
        """Export DPO training data.

        Args:
            output_path: Path to output file.
            examples: List of TrainingExample objects with preference data.

        Returns:
            Number of examples exported.
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        count = 0
        with open(output_path, "w") as f:
            for example in examples:
                pref_data = example.to_preference_format()
                if pref_data:
                    f.write(json.dumps(pref_data) + "\n")
                    count += 1

        logger.info(f"Exported {count} DPO examples to {output_path}")
        return count

    def export_full(
        self,
        output_path: str | Path,
        examples: list[TrainingExample],
    ) -> int:
        """Export full training example data (including metadata).

        Args:
            output_path: Path to output file.
            examples: List of TrainingExample objects.

        Returns:
            Number of examples exported.
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w") as f:
            for example in examples:
                f.write(json.dumps(example.to_dict()) + "\n")

        logger.info(f"Exported {len(examples)} full examples to {output_path}")
        return len(examples)

    def create_train_val_split(
        self,
        examples: list[TrainingExample],
        val_ratio: float = 0.1,
        seed: int = 42,
    ) -> tuple[list[TrainingExample], list[TrainingExample]]:
        """Split examples into training and validation sets.

        Args:
            examples: List of TrainingExample objects.
            val_ratio: Fraction for validation set.
            seed: Random seed for reproducibility.

        Returns:
            Tuple of (train_examples, val_examples).
        """
        import random

        random.seed(seed)
        shuffled = examples.copy()
        random.shuffle(shuffled)

        split_idx = int(len(shuffled) * (1 - val_ratio))
        train = shuffled[:split_idx]
        val = shuffled[split_idx:]

        logger.info(f"Split {len(examples)} examples: {len(train)} train, {len(val)} val")
        return train, val

    def get_statistics(
        self,
        examples: list[TrainingExample],
    ) -> dict[str, Any]:
        """Get statistics about the training data.

        Args:
            examples: List of TrainingExample objects.

        Returns:
            Dictionary of statistics.
        """
        stats = {
            "total_examples": len(examples),
            "expert_validated": sum(1 for e in examples if e.expert_validated),
            "with_context": sum(1 for e in examples if e.context),
            "with_preference": sum(1 for e in examples if e.preferred_response),
            "avg_accuracy_score": 0.0,
            "avg_conservatism_score": 0.0,
            "tag_distribution": {},
        }

        # Calculate averages for validated examples
        validated = [e for e in examples if e.expert_validated]
        if validated:
            stats["avg_accuracy_score"] = sum(e.accuracy_score for e in validated) / len(validated)
            stats["avg_conservatism_score"] = sum(e.conservatism_score for e in validated) / len(validated)

        # Tag distribution
        for example in examples:
            for tag in example.tags:
                stats["tag_distribution"][tag] = stats["tag_distribution"].get(tag, 0) + 1

        return stats
