"""Trax-specific repository base classes built on AI Assistant Library."""

# Simplified base classes for repository pattern
from abc import ABC
from datetime import datetime
from typing import Any, Dict, List, Optional
from uuid import UUID, uuid4

from sqlalchemy import Column, DateTime, Float, Integer, String, Text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.dialects.postgresql import UUID as PGUUID
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session

Base = declarative_base()


class BaseModel(Base):
    """Base model with common fields."""

    __abstract__ = True

    def to_dict(self) -> Dict[str, Any]:
        """Convert model to dictionary."""
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}


class TimestampedModel(BaseModel):
    """Model with timestamp fields."""

    __abstract__ = True

    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)


class VersionedModel(TimestampedModel):
    """Model with version tracking."""

    __abstract__ = True

    version = Column(Integer, default=1, nullable=False)


class BaseRepository(ABC):
    """Base repository for database operations."""

    def __init__(self, session, model_class):
        self.session = session
        self.model_class = model_class

    async def create(self, data: Dict[str, Any]):
        """Create a new entity."""
        entity = self.model_class(**data)
        self.session.add(entity)
        self.session.commit()
        return entity

    async def find_by_id(self, entity_id):
        """Find entity by ID."""
        return self.session.query(self.model_class).filter(self.model_class.id == entity_id).first()

    async def find_all(self, filters: Optional[Dict[str, Any]] = None):
        """Find all entities matching filters."""
        query = self.session.query(self.model_class)
        if filters:
            for key, value in filters.items():
                if hasattr(self.model_class, key):
                    query = query.filter(getattr(self.model_class, key) == value)
        return query.all()

    async def update(self, entity_id, data: Dict[str, Any]):
        """Update an entity."""
        entity = await self.find_by_id(entity_id)
        if entity:
            for key, value in data.items():
                if hasattr(entity, key):
                    setattr(entity, key, value)
            self.session.commit()
        return entity

    async def delete(self, entity_id):
        """Delete an entity."""
        entity = await self.find_by_id(entity_id)
        if entity:
            self.session.delete(entity)
            self.session.commit()
            return True
        return False


class TimestampedRepository(BaseRepository):
    """Repository for timestamped models."""

    pass


class MediaFile(TimestampedModel):
    """Media file model with PostgreSQL JSONB support."""

    __tablename__ = "media_files"

    id = Column(PGUUID(as_uuid=True), primary_key=True, default=uuid4)
    source_path = Column(Text, nullable=False)
    local_path = Column(Text)
    file_hash = Column(String(64), unique=True, index=True)
    file_size = Column(Float)
    mime_type = Column(String(100))
    duration_seconds = Column(Float)
    file_metadata = Column(JSONB, default=dict)  # PostgreSQL JSONB for flexible metadata

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            "id": str(self.id),
            "source_path": self.source_path,
            "local_path": self.local_path,
            "file_hash": self.file_hash,
            "file_size": self.file_size,
            "mime_type": self.mime_type,
            "duration_seconds": self.duration_seconds,
            "file_metadata": self.file_metadata,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }


class Transcript(VersionedModel):
    """Transcript model with versioning and JSONB content."""

    __tablename__ = "transcripts"

    id = Column(PGUUID(as_uuid=True), primary_key=True, default=uuid4)
    media_file_id = Column(PGUUID(as_uuid=True), nullable=False, index=True)

    # Version tracking for v1, v2, v3, v4
    pipeline_version = Column(String(10), nullable=False, default="v1")

    # JSONB columns for flexible content storage
    raw_content = Column(JSONB, nullable=False)
    enhanced_content = Column(JSONB)
    multipass_content = Column(JSONB)
    diarized_content = Column(JSONB)

    # Metadata
    confidence_scores = Column(JSONB)
    speaker_profiles = Column(JSONB)
    processing_time = Column(Float)
    model_config = Column(JSONB)

    # Parent transcript for version tracking
    parent_transcript_id = Column(PGUUID(as_uuid=True), index=True)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            "id": str(self.id),
            "media_file_id": str(self.media_file_id),
            "pipeline_version": self.pipeline_version,
            "raw_content": self.raw_content,
            "enhanced_content": self.enhanced_content,
            "multipass_content": self.multipass_content,
            "diarized_content": self.diarized_content,
            "confidence_scores": self.confidence_scores,
            "speaker_profiles": self.speaker_profiles,
            "processing_time": self.processing_time,
            "model_config": self.model_config,
            "parent_transcript_id": (
                str(self.parent_transcript_id) if self.parent_transcript_id else None
            ),
            "version": self.version,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }


class ProcessingJob(TimestampedModel):
    """Processing job model for tracking batch operations."""

    __tablename__ = "processing_jobs"

    id = Column(PGUUID(as_uuid=True), primary_key=True, default=uuid4)
    job_type = Column(String(50), nullable=False, index=True)
    status = Column(String(20), nullable=False, index=True)
    media_file_id = Column(PGUUID(as_uuid=True), index=True)

    # Job configuration and results
    config = Column(JSONB, default=dict)
    result = Column(JSONB)
    error = Column(Text)

    # Timing
    started_at = Column(DateTime)
    completed_at = Column(DateTime)

    # Progress tracking
    total_items = Column(Float, default=0)
    processed_items = Column(Float, default=0)
    successful_items = Column(Float, default=0)
    failed_items = Column(Float, default=0)

    @property
    def progress_percentage(self) -> float:
        """Calculate progress percentage."""
        if self.total_items == 0:
            return 0.0
        return (self.processed_items / self.total_items) * 100

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            "id": str(self.id),
            "job_type": self.job_type,
            "status": self.status,
            "media_file_id": str(self.media_file_id) if self.media_file_id else None,
            "config": self.config,
            "result": self.result,
            "error": self.error,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "progress_percentage": self.progress_percentage,
            "total_items": self.total_items,
            "processed_items": self.processed_items,
            "successful_items": self.successful_items,
            "failed_items": self.failed_items,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }


class MediaFileRepository(TimestampedRepository):
    """Repository for media file operations."""

    def __init__(self, session: Session):
        """Initialize media file repository.

        Args:
            session: SQLAlchemy database session

        """
        super().__init__(session, MediaFile)

    async def find_by_hash(self, file_hash: str) -> Optional[MediaFile]:
        """Find media file by hash.

        Args:
            file_hash: SHA256 hash of the file

        Returns:
            MediaFile if found, None otherwise

        """
        return self.session.query(MediaFile).filter(MediaFile.file_hash == file_hash).first()

    async def find_by_source_path(self, source_path: str) -> Optional[MediaFile]:
        """Find media file by source path.

        Args:
            source_path: Original source path

        Returns:
            MediaFile if found, None otherwise

        """
        return self.session.query(MediaFile).filter(MediaFile.source_path == source_path).first()

    async def update_metadata(
        self, media_file_id: UUID, metadata: Dict[str, Any]
    ) -> Optional[MediaFile]:
        """Update media file metadata.

        Args:
            media_file_id: Media file ID
            metadata: New metadata to merge

        Returns:
            Updated MediaFile if found

        """
        media_file = await self.find_by_id(media_file_id)
        if media_file:
            # Merge metadata
            current_metadata = media_file.metadata or {}
            current_metadata.update(metadata)
            media_file.metadata = current_metadata
            self.session.commit()
        return media_file


class TranscriptRepository(TimestampedRepository):
    """Repository for transcript operations with version support."""

    def __init__(self, session: Session):
        """Initialize transcript repository.

        Args:
            session: SQLAlchemy database session

        """
        super().__init__(session, Transcript)

    async def find_by_media_file(
        self, media_file_id: UUID, version: Optional[str] = None
    ) -> List[Transcript]:
        """Find transcripts for a media file.

        Args:
            media_file_id: Media file ID
            version: Optional pipeline version filter

        Returns:
            List of transcripts

        """
        query = self.session.query(Transcript).filter(Transcript.media_file_id == media_file_id)

        if version:
            query = query.filter(Transcript.pipeline_version == version)

        return query.all()

    async def find_latest_version(self, media_file_id: UUID) -> Optional[Transcript]:
        """Find the latest transcript version for a media file.

        Args:
            media_file_id: Media file ID

        Returns:
            Latest transcript if found

        """
        # Order by pipeline version (v4 > v3 > v2 > v1) and created_at
        return (
            self.session.query(Transcript)
            .filter(Transcript.media_file_id == media_file_id)
            .order_by(Transcript.pipeline_version.desc(), Transcript.created_at.desc())
            .first()
        )

    async def create_new_version(
        self,
        media_file_id: UUID,
        content: Dict[str, Any],
        version: str,
        parent_id: Optional[UUID] = None,
    ) -> Transcript:
        """Create a new transcript version.

        Args:
            media_file_id: Media file ID
            content: Transcript content
            version: Pipeline version
            parent_id: Parent transcript ID for version tracking

        Returns:
            Created transcript

        """
        transcript_data = {
            "media_file_id": media_file_id,
            "pipeline_version": version,
            "raw_content": content.get("raw_content", content),
            "enhanced_content": content.get("enhanced_content"),
            "multipass_content": content.get("multipass_content"),
            "diarized_content": content.get("diarized_content"),
            "confidence_scores": content.get("confidence_scores"),
            "speaker_profiles": content.get("speaker_profiles"),
            "processing_time": content.get("processing_time"),
            "model_config": content.get("model_config"),
            "parent_transcript_id": parent_id,
        }

        return await self.create(transcript_data)


class ProcessingJobRepository(TimestampedRepository):
    """Repository for processing job operations."""

    def __init__(self, session: Session):
        """Initialize processing job repository.

        Args:
            session: SQLAlchemy database session

        """
        super().__init__(session, ProcessingJob)

    async def find_active_jobs(self) -> List[ProcessingJob]:
        """Find all active processing jobs.

        Returns:
            List of active jobs

        """
        return (
            self.session.query(ProcessingJob)
            .filter(ProcessingJob.status.in_(["pending", "processing"]))
            .all()
        )

    async def update_progress(
        self, job_id: UUID, processed: int, successful: int, failed: int
    ) -> Optional[ProcessingJob]:
        """Update job progress.

        Args:
            job_id: Job ID
            processed: Number of processed items
            successful: Number of successful items
            failed: Number of failed items

        Returns:
            Updated job if found

        """
        job = await self.find_by_id(job_id)
        if job:
            job.processed_items = processed
            job.successful_items = successful
            job.failed_items = failed

            # Update status if complete
            if processed >= job.total_items:
                job.status = "completed" if failed == 0 else "partial"
                job.completed_at = datetime.utcnow()

            self.session.commit()
        return job

    async def mark_failed(self, job_id: UUID, error: str) -> Optional[ProcessingJob]:
        """Mark job as failed.

        Args:
            job_id: Job ID
            error: Error message

        Returns:
            Updated job if found

        """
        job = await self.find_by_id(job_id)
        if job:
            job.status = "failed"
            job.error = error
            job.completed_at = datetime.utcnow()
            self.session.commit()
        return job
