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

from typing import TYPE_CHECKING

from sqlalchemy import JSON, CheckConstraint, Column, Float, ForeignKey, Integer, String
from sqlalchemy.orm import relationship

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

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


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

    __tablename__ = "transcripts"

    id = Column(Integer, autoincrement=True, primary_key=True)
    model = Column(
        String, nullable=False, comment="Name of model used to generate transcript"
    )
    media_id = Column(
        MediaIdSqlType,
        ForeignKey("video.id"),
        nullable=False,
        comment="Foreign key to video table",
    )
    transcript = Column(String, comment="Full text transcript of media")
    segments = Column(JSON, comment="Segments of the transcript")
    language = Column(
        String, comment="Language of the transcript as predicted by model"
    )
    language_confidence = Column(
        Float,
        CheckConstraint(
            "0 <= language_confidence <= 1", name="language_confidence_value_range"
        ),
        comment="Confidence score of language prediction",
    )

    video = relationship("VideoEntity", back_populates="transcripts", uselist=False)

    @classmethod
    def from_asr_output(
        cls,
        model_name: str,
        media_id: MediaId,
        info: AsrTranscriptionInfo,
        transcription: AsrTranscription,
        segments: AsrSegments,
    ) -> TranscriptEntity:
        """Converts an AsrTranscriptionInfo and AsrTranscription to a single Transcript entity."""
        return TranscriptEntity(
            model=model_name,
            media_id=media_id,
            language=info.language,
            language_confidence=info.language_confidence,
            transcript=transcription.text,
            segments=[s.model_dump() for s in segments],
        )
