"""
Training Readiness Assessor

Evaluates whether collected data is sufficient for custom model training.
"""

import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

from sage.data_collection.schemas import (
    ExpertFeedback,
    FeedbackRating,
    LoggedInteraction,
)
from sage.analysis.gap_analyzer import GapReport, PerformanceGapAnalyzer

logger = logging.getLogger(__name__)


@dataclass
class DataQualityMetrics:
    """Metrics about training data quality."""
    total_interactions: int = 0
    expert_validated_count: int = 0
    high_quality_count: int = 0  # Score >= 4
    with_corrections_count: int = 0
    topic_coverage: dict[str, int] = field(default_factory=dict)
    query_type_coverage: dict[str, int] = field(default_factory=dict)


@dataclass
class TrainingReadinessThresholds:
    """Thresholds for training readiness decision."""
    min_total_interactions: int = 1000
    min_expert_validated: int = 200
    min_high_quality: int = 150
    min_preference_pairs: int = 50
    min_topic_coverage: int = 5  # Topics with >= 50 examples
    min_query_type_coverage: int = 4  # Query types with >= 100 examples


@dataclass
class TrainingReadinessReport:
    """Report on readiness for custom model training."""

    # Decision
    ready_for_training: bool = False
    confidence: float = 0.0  # 0-1 confidence in decision

    # Data quantity
    total_interactions: int = 0
    expert_validated: int = 0
    high_quality_examples: int = 0
    preference_pairs: int = 0

    # Data quality
    quality_metrics: DataQualityMetrics = field(default_factory=DataQualityMetrics)

    # Coverage
    topic_coverage_score: float = 0.0
    query_type_coverage_score: float = 0.0
    coverage_gaps: list[str] = field(default_factory=list)

    # Gap analysis summary
    performance_gap_summary: str = ""
    can_address_with_prompts: bool = True

    # Thresholds comparison
    thresholds_met: dict[str, bool] = field(default_factory=dict)

    # Recommendations
    recommendation: str = ""
    next_steps: list[str] = field(default_factory=list)

    # Metadata
    assessment_timestamp: datetime = field(default_factory=datetime.now)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "decision": {
                "ready_for_training": self.ready_for_training,
                "confidence": self.confidence,
                "recommendation": self.recommendation,
            },
            "data_quantity": {
                "total_interactions": self.total_interactions,
                "expert_validated": self.expert_validated,
                "high_quality_examples": self.high_quality_examples,
                "preference_pairs": self.preference_pairs,
            },
            "data_quality": {
                "topic_coverage": self.quality_metrics.topic_coverage,
                "query_type_coverage": self.quality_metrics.query_type_coverage,
            },
            "coverage": {
                "topic_coverage_score": self.topic_coverage_score,
                "query_type_coverage_score": self.query_type_coverage_score,
                "gaps": self.coverage_gaps,
            },
            "thresholds_met": self.thresholds_met,
            "next_steps": self.next_steps,
            "performance_analysis": {
                "gap_summary": self.performance_gap_summary,
                "can_address_with_prompts": self.can_address_with_prompts,
            },
            "metadata": {
                "assessment_timestamp": self.assessment_timestamp.isoformat(),
            },
        }


class TrainingReadinessAssessor:
    """Assesses readiness for custom model training.

    Evaluates:
    1. Data quantity (enough examples?)
    2. Data quality (expert validated?)
    3. Topic coverage (all important areas?)
    4. Performance gaps (addressable with training?)

    Example:
        assessor = TrainingReadinessAssessor()

        report = assessor.assess(
            interactions=logged_interactions,
            feedback_map=expert_feedback,
        )

        if report.ready_for_training:
            print("Ready for training!")
        else:
            print(f"Not ready: {report.recommendation}")
            print(f"Next steps: {report.next_steps}")
    """

    def __init__(
        self,
        thresholds: TrainingReadinessThresholds | None = None,
    ):
        """Initialize assessor.

        Args:
            thresholds: Custom thresholds for readiness decision.
        """
        self.thresholds = thresholds or TrainingReadinessThresholds()
        self.gap_analyzer = PerformanceGapAnalyzer()

    def assess(
        self,
        interactions: list[LoggedInteraction],
        feedback_map: dict[str, ExpertFeedback] | None = None,
        gap_report: GapReport | None = None,
    ) -> TrainingReadinessReport:
        """Assess training readiness.

        Args:
            interactions: List of logged interactions.
            feedback_map: Map of interaction_id -> expert feedback.
            gap_report: Optional pre-computed gap report.

        Returns:
            TrainingReadinessReport with decision and recommendations.
        """
        feedback_map = feedback_map or {}
        report = TrainingReadinessReport()

        # Calculate data quantity metrics
        self._assess_data_quantity(interactions, feedback_map, report)

        # Calculate data quality metrics
        self._assess_data_quality(interactions, feedback_map, report)

        # Assess coverage
        self._assess_coverage(interactions, report)

        # Analyze performance gaps
        if gap_report is None:
            gap_report = self.gap_analyzer.analyze(interactions, feedback_map)
        self._assess_performance_gaps(gap_report, report)

        # Check thresholds
        self._check_thresholds(report)

        # Make decision
        self._make_decision(report)

        logger.info(
            f"Training readiness assessment: {'READY' if report.ready_for_training else 'NOT READY'} "
            f"(confidence: {report.confidence:.0%})"
        )

        return report

    def _assess_data_quantity(
        self,
        interactions: list[LoggedInteraction],
        feedback_map: dict[str, ExpertFeedback],
        report: TrainingReadinessReport,
    ) -> None:
        """Assess data quantity metrics."""
        report.total_interactions = len(interactions)
        report.expert_validated = len(feedback_map)

        # Count high-quality examples (score >= 4)
        report.high_quality_examples = sum(
            1 for f in feedback_map.values()
            if f.accuracy_score.value >= 4
        )

        # Count preference pairs (have corrections)
        report.preference_pairs = sum(
            1 for f in feedback_map.values()
            if f.corrected_response is not None
        )

        report.quality_metrics.total_interactions = len(interactions)
        report.quality_metrics.expert_validated_count = report.expert_validated
        report.quality_metrics.high_quality_count = report.high_quality_examples
        report.quality_metrics.with_corrections_count = report.preference_pairs

    def _assess_data_quality(
        self,
        interactions: list[LoggedInteraction],
        feedback_map: dict[str, ExpertFeedback],
        report: TrainingReadinessReport,
    ) -> None:
        """Assess data quality metrics."""
        # Topic coverage
        topic_counts: dict[str, int] = {}
        for interaction in interactions:
            topics = self._extract_topics(interaction.query)
            for topic in topics:
                topic_counts[topic] = topic_counts.get(topic, 0) + 1

        report.quality_metrics.topic_coverage = topic_counts

        # Query type coverage
        type_counts: dict[str, int] = {}
        for interaction in interactions:
            qtype = interaction.query_type
            type_counts[qtype] = type_counts.get(qtype, 0) + 1

        report.quality_metrics.query_type_coverage = type_counts

    def _assess_coverage(
        self,
        interactions: list[LoggedInteraction],
        report: TrainingReadinessReport,
    ) -> None:
        """Assess topic and query type coverage."""
        # Topic coverage score (how many topics have >= 50 examples)
        topics_with_coverage = sum(
            1 for count in report.quality_metrics.topic_coverage.values()
            if count >= 50
        )
        total_topics = len(report.quality_metrics.topic_coverage)
        report.topic_coverage_score = topics_with_coverage / max(total_topics, 1)

        # Query type coverage score
        types_with_coverage = sum(
            1 for count in report.quality_metrics.query_type_coverage.values()
            if count >= 100
        )
        total_types = len(report.quality_metrics.query_type_coverage)
        report.query_type_coverage_score = types_with_coverage / max(total_types, 1)

        # Identify coverage gaps
        gaps = []
        for topic, count in report.quality_metrics.topic_coverage.items():
            if count < 20:
                gaps.append(f"{topic} (only {count} examples)")

        for qtype, count in report.quality_metrics.query_type_coverage.items():
            if count < 50:
                gaps.append(f"{qtype} queries (only {count} examples)")

        report.coverage_gaps = gaps

    def _assess_performance_gaps(
        self,
        gap_report: GapReport,
        report: TrainingReadinessReport,
    ) -> None:
        """Assess whether performance gaps require training or can be fixed with prompts."""
        # Summarize gaps
        gap_summary_parts = []

        if gap_report.topic_gaps:
            top_gaps = [g.topic for g in gap_report.topic_gaps[:3]]
            gap_summary_parts.append(f"Weak topics: {', '.join(top_gaps)}")

        if gap_report.factual_error_rate > 0.05:
            gap_summary_parts.append(
                f"Factual error rate: {gap_report.factual_error_rate:.1%}"
            )

        if gap_report.false_positive_rate > 0.1:
            gap_summary_parts.append(
                f"False positive rate: {gap_report.false_positive_rate:.1%}"
            )

        report.performance_gap_summary = "; ".join(gap_summary_parts) or "No significant gaps"

        # Determine if gaps can be addressed with prompts
        # If gaps are mostly about abstention tuning, they can be fixed with prompts
        # If gaps are about factual accuracy, training is needed
        report.can_address_with_prompts = (
            gap_report.factual_error_rate < 0.05 and
            len(gap_report.topic_gaps) < 3 and
            gap_report.average_accuracy_score >= 3.5
        )

    def _check_thresholds(self, report: TrainingReadinessReport) -> None:
        """Check which thresholds are met."""
        t = self.thresholds

        report.thresholds_met = {
            "total_interactions": report.total_interactions >= t.min_total_interactions,
            "expert_validated": report.expert_validated >= t.min_expert_validated,
            "high_quality": report.high_quality_examples >= t.min_high_quality,
            "preference_pairs": report.preference_pairs >= t.min_preference_pairs,
            "topic_coverage": sum(
                1 for c in report.quality_metrics.topic_coverage.values() if c >= 50
            ) >= t.min_topic_coverage,
            "query_type_coverage": sum(
                1 for c in report.quality_metrics.query_type_coverage.values() if c >= 100
            ) >= t.min_query_type_coverage,
        }

    def _make_decision(self, report: TrainingReadinessReport) -> None:
        """Make the training readiness decision."""
        # Calculate confidence based on thresholds met
        thresholds_met = sum(report.thresholds_met.values())
        total_thresholds = len(report.thresholds_met)
        report.confidence = thresholds_met / total_thresholds

        # Decision: ready if all critical thresholds met OR
        # if most met and gaps can't be addressed with prompts
        critical_met = (
            report.thresholds_met.get("total_interactions", False) and
            report.thresholds_met.get("expert_validated", False) and
            report.thresholds_met.get("high_quality", False)
        )

        if critical_met and report.confidence >= 0.7:
            report.ready_for_training = True
            report.recommendation = (
                "Data quantity and quality thresholds met. "
                "Proceed with custom model training to address identified gaps."
            )
        elif report.can_address_with_prompts and report.confidence < 0.5:
            report.ready_for_training = False
            report.recommendation = (
                "Current performance gaps can likely be addressed with prompt engineering. "
                "Continue collecting data before investing in custom training."
            )
        else:
            report.ready_for_training = False
            report.recommendation = (
                "Insufficient training data. "
                f"Met {thresholds_met}/{total_thresholds} thresholds."
            )

        # Generate next steps
        self._generate_next_steps(report)

    def _generate_next_steps(self, report: TrainingReadinessReport) -> None:
        """Generate recommended next steps."""
        next_steps = []

        if not report.thresholds_met.get("total_interactions", True):
            gap = self.thresholds.min_total_interactions - report.total_interactions
            next_steps.append(f"Collect {gap:,} more interactions")

        if not report.thresholds_met.get("expert_validated", True):
            gap = self.thresholds.min_expert_validated - report.expert_validated
            next_steps.append(f"Get expert reviews for {gap} more interactions")

        if not report.thresholds_met.get("preference_pairs", True):
            gap = self.thresholds.min_preference_pairs - report.preference_pairs
            next_steps.append(f"Collect {gap} more corrections for preference learning")

        if report.coverage_gaps:
            next_steps.append(f"Address coverage gaps: {', '.join(report.coverage_gaps[:3])}")

        if report.ready_for_training:
            next_steps = [
                "Prepare training data export",
                "Select base model (Llama 3.1, Mistral, etc.)",
                "Set up training infrastructure",
                "Plan SFT + DPO training pipeline",
            ]

        report.next_steps = next_steps

    def _extract_topics(self, query: str) -> list[str]:
        """Extract topics from query for coverage analysis."""
        # Reuse gap analyzer's topic extraction
        return self.gap_analyzer._extract_topics(query)

    def get_data_requirements(self) -> dict[str, int]:
        """Get minimum data requirements for training readiness.

        Returns:
            Dictionary of requirement name -> minimum count.
        """
        return {
            "total_interactions": self.thresholds.min_total_interactions,
            "expert_validated": self.thresholds.min_expert_validated,
            "high_quality_examples": self.thresholds.min_high_quality,
            "preference_pairs": self.thresholds.min_preference_pairs,
        }
