# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from abc import ABC
from enum import Enum

from pydantic import BaseModel, Field
from rich.panel import Panel
from rich.table import Column, Table
from typing_extensions import TypeAlias

from data_designer.config.analysis.column_statistics import (
    CategoricalDistribution,
    CategoricalHistogramData,
    ColumnDistributionType,
    MissingValue,
    NumericalDistribution,
)
from data_designer.config.analysis.utils.reporting import TITLE_STYLE, create_judge_score_summary_table
from data_designer.config.base import ConfigBase
from data_designer.config.utils.visualization import ColorPalette


class ColumnProfilerType(str, Enum):
    JUDGE_SCORE = "judge-score"


class ColumnProfilerResults(BaseModel, ABC):
    """Abstract base class for column profiler results.

    Stores results from column profiling operations. Subclasses hold profiler-specific
    analysis results and provide methods for generating formatted report sections for display.
    """

    def create_report_section(self) -> Panel:
        """Creates a Rich Panel containing the formatted profiler results for display.

        Returns:
            A Rich Panel containing the formatted profiler results. Default implementation
            returns a "Not Implemented" message; subclasses should override to provide
            specific formatting.
        """
        return Panel(
            f"Report section generation not implemented for '{self.__class__.__name__}'.",
            title="Not Implemented",
            border_style=f"bold {ColorPalette.YELLOW.value}",
            padding=(1, 2),
        )


class JudgeScoreProfilerConfig(ConfigBase):
    """Configuration for the LLM-as-a-judge score profiler.

    Attributes:
        model_alias: Alias of the LLM model to use for generating score distribution summaries.
            Must match a model alias defined in the Data Designer configuration.
        summary_score_sample_size: Number of score samples to include when prompting the LLM
            to generate summaries. Larger sample sizes provide more context but increase
            token usage. Must be at least 1. Defaults to 20.
    """

    model_alias: str
    summary_score_sample_size: int | None = Field(default=20, ge=1)


class JudgeScoreSample(BaseModel):
    """Container for a single judge score and its associated reasoning.

    Stores a paired score-reasoning sample extracted from an LLM-as-a-judge column.
    Used when generating summaries to provide the LLM with examples of scoring patterns.

    Attributes:
        score: The score value assigned by the judge. Can be numeric (int) or categorical (str).
        reasoning: The reasoning or explanation provided by the judge for this score.
    """

    score: int | str
    reasoning: str


class JudgeScoreDistributions(BaseModel):
    """Container for computed distributions across all judge score dimensions.

    Stores the complete distribution analysis for all score dimensions in an LLM-as-a-judge
    column. Each score dimension (e.g., "relevance", "fluency") has its own distribution
    computed from the generated data.

    Attributes:
        scores: Mapping of each score dimension name to its list of score values.
        reasoning: Mapping of each score dimension name to its list of reasoning texts.
        distribution_types: Mapping of each score dimension name to its classification.
        distributions: Mapping of each score dimension name to its computed distribution statistics.
        histograms: Mapping of each score dimension name to its histogram data.
    """

    scores: dict[str, list[int | str]]
    reasoning: dict[str, list[str]]
    distribution_types: dict[str, ColumnDistributionType]
    distributions: dict[str, CategoricalDistribution | NumericalDistribution | MissingValue]
    histograms: dict[str, CategoricalHistogramData | MissingValue]


class JudgeScoreSummary(BaseModel):
    """Container for an LLM-generated summary of a judge score dimension.

    Stores the natural language summary and sample data for a single score dimension
    generated by the judge score profiler. The summary is created by an LLM analyzing
    the distribution and patterns in the score-reasoning pairs.

    Attributes:
        score_name: Name of the score dimension being summarized (e.g., "relevance", "fluency").
        summary: LLM-generated natural language summary describing the scoring patterns,
            distribution characteristics, and notable trends for this score dimension.
        score_samples: List of score-reasoning pairs that were used to generate the summary.
            These are the examples of the scoring behavior that were used to generate the summary.
    """

    score_name: str
    summary: str
    score_samples: list[JudgeScoreSample]


class JudgeScoreProfilerResults(ColumnProfilerResults):
    """Container for complete judge score profiler analysis results.

    Attributes:
        column_name: Name of the judge column that was profiled.
        summaries: Mapping of each score dimension name to its LLM-generated summary.
        score_distributions: Complete distribution analysis across all score dimensions.
    """

    column_name: str
    summaries: dict[str, JudgeScoreSummary]
    score_distributions: JudgeScoreDistributions | MissingValue

    def create_report_section(self) -> Panel:
        layout = Table.grid(Column(), expand=True, padding=(2, 0))

        for score_name in self.summaries.keys():
            layout.add_row(
                create_judge_score_summary_table(
                    score_name=score_name,
                    histogram=self.score_distributions.histograms[score_name],
                    summary=self.summaries[score_name].summary,
                )
            )

        return Panel(
            layout,
            title=f"[{TITLE_STYLE}]LLM-as-a-Judge Score Profile: '{self.column_name}'[/{TITLE_STYLE}]",
            padding=(1, 2),
        )


ColumnProfilerConfigT: TypeAlias = JudgeScoreProfilerConfig

ColumnProfilerResultsT: TypeAlias = JudgeScoreProfilerResults
