from __future__ import annotations  # Let classes use themselves in type annotations

from typing import TYPE_CHECKING

from sqlalchemy import JSON, CheckConstraint
from sqlalchemy.orm import Mapped, mapped_column

from aana.storage.models.base import BaseEntity, TimeStampEntity

if TYPE_CHECKING:
    from aana.core.models.asr import (
        AsrSegments,
        AsrTranscription,
        AsrTranscriptionInfo,
    )


class TranscriptEntity(BaseEntity, TimeStampEntity):
    """ORM class for media transcripts generated by a model.

    Attributes:
        id (int): Unique identifier for the transcript.
        model (str): Name of the model used to generate the transcript.
        transcript (str): Full text transcript of the media.
        segments (dict): Segments of the transcript.
        language (str): Language of the transcript as predicted by the model.
        language_confidence (float): Confidence score of language prediction.
        transcript_type (str): The type of transcript (populated automatically by ORM based on `polymorphic_identity` of subclass).
    """

    __tablename__ = "transcript"

    id: Mapped[int] = mapped_column(autoincrement=True, primary_key=True)
    model: Mapped[str] = mapped_column(
        nullable=False, comment="Name of model used to generate transcript"
    )
    transcript: Mapped[str] = mapped_column(comment="Full text transcript of media")
    segments: Mapped[dict] = mapped_column(JSON, comment="Segments of the transcript")
    language: Mapped[str] = mapped_column(
        comment="Language of the transcript as predicted by model"
    )
    language_confidence: Mapped[float] = mapped_column(
        CheckConstraint(
            "0 <= language_confidence <= 1", name="language_confidence_value_range"
        ),
        comment="Confidence score of language prediction",
    )
    transcript_type: Mapped[str] = mapped_column(comment="The type of transcript")

    __mapper_args__ = {  # noqa: RUF012
        "polymorphic_identity": "transcript",
        "polymorphic_on": "transcript_type",
    }

    @classmethod
    def from_asr_output(
        cls,
        model_name: str,
        info: AsrTranscriptionInfo,
        transcription: AsrTranscription,
        segments: AsrSegments,
    ) -> TranscriptEntity:
        """Converts an AsrTranscriptionInfo and AsrTranscription to a single Transcript entity.

        Args:
            model_name (str): Name of the model used to generate the transcript.
            info (AsrTranscriptionInfo): Information about the transcription.
            transcription (AsrTranscription): The full transcription.
            segments (AsrSegments): Segments of the transcription.

        Returns:
            TranscriptEntity: A new instance of the TranscriptEntity class.
        """
        return TranscriptEntity(
            model=model_name,
            language=info.language,
            language_confidence=info.language_confidence,
            transcript=transcription.text,
            segments=[s.model_dump() for s in segments],
        )
