"""
Document Repository

Handles document operations following Single Responsibility Principle.
"""

import logging
from datetime import datetime, timezone
from typing import Any
from uuid import UUID

from sqlalchemy import func
from sqlmodel import select

from backend.data.database import get_session_context
from backend.data.admin_models import AdminDocument, AdminAnnotation
from backend.data.repositories.base import BaseRepository

logger = logging.getLogger(__name__)


class DocumentRepository(BaseRepository[AdminDocument]):
    """Repository for document management.

    Handles:
    - Document CRUD operations
    - Document status management
    - Document filtering and pagination
    - Document category management
    """

    def create(
        self,
        filename: str,
        file_size: int,
        content_type: str,
        file_path: str,
        page_count: int = 1,
        upload_source: str = "ui",
        csv_field_values: dict[str, Any] | None = None,
        group_key: str | None = None,
        category: str = "invoice",
        admin_token: str | None = None,
    ) -> str:
        """Create a new document record.

        Args:
            filename: Original filename
            file_size: File size in bytes
            content_type: MIME type
            file_path: Storage path
            page_count: Number of pages
            upload_source: Upload source (ui/api)
            csv_field_values: CSV field values for reference
            group_key: User-defined grouping key
            category: Document category
            admin_token: Deprecated, kept for compatibility

        Returns:
            Document ID as string
        """
        with get_session_context() as session:
            document = AdminDocument(
                filename=filename,
                file_size=file_size,
                content_type=content_type,
                file_path=file_path,
                page_count=page_count,
                upload_source=upload_source,
                csv_field_values=csv_field_values,
                group_key=group_key,
                category=category,
            )
            session.add(document)
            session.flush()
            return str(document.document_id)

    def get(self, document_id: str) -> AdminDocument | None:
        """Get a document by ID.

        Args:
            document_id: Document UUID as string

        Returns:
            AdminDocument if found, None otherwise
        """
        with get_session_context() as session:
            result = session.get(AdminDocument, UUID(document_id))
            if result:
                session.expunge(result)
            return result

    def get_by_token(
        self,
        document_id: str,
        admin_token: str | None = None,
    ) -> AdminDocument | None:
        """Get a document by ID. Token parameter is deprecated."""
        return self.get(document_id)

    def get_paginated(
        self,
        admin_token: str | None = None,
        status: str | None = None,
        upload_source: str | None = None,
        has_annotations: bool | None = None,
        auto_label_status: str | None = None,
        batch_id: str | None = None,
        category: str | None = None,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[AdminDocument], int]:
        """Get paginated documents with optional filters.

        Args:
            admin_token: Deprecated, kept for compatibility
            status: Filter by status
            upload_source: Filter by upload source
            has_annotations: Filter by annotation presence
            auto_label_status: Filter by auto-label status
            batch_id: Filter by batch ID
            category: Filter by category
            limit: Page size
            offset: Pagination offset

        Returns:
            Tuple of (documents, total_count)
        """
        with get_session_context() as session:
            where_clauses = []

            if status:
                where_clauses.append(AdminDocument.status == status)
            if upload_source:
                where_clauses.append(AdminDocument.upload_source == upload_source)
            if auto_label_status:
                where_clauses.append(AdminDocument.auto_label_status == auto_label_status)
            if batch_id:
                where_clauses.append(AdminDocument.batch_id == UUID(batch_id))
            if category:
                where_clauses.append(AdminDocument.category == category)

            count_stmt = select(func.count()).select_from(AdminDocument)
            if where_clauses:
                count_stmt = count_stmt.where(*where_clauses)

            if has_annotations is not None:
                if has_annotations:
                    count_stmt = (
                        count_stmt
                        .join(AdminAnnotation, AdminAnnotation.document_id == AdminDocument.document_id)
                        .group_by(AdminDocument.document_id)
                    )
                else:
                    count_stmt = (
                        count_stmt
                        .outerjoin(AdminAnnotation, AdminAnnotation.document_id == AdminDocument.document_id)
                        .where(AdminAnnotation.annotation_id.is_(None))
                    )

            total = session.exec(count_stmt).one()

            statement = select(AdminDocument)
            if where_clauses:
                statement = statement.where(*where_clauses)

            if has_annotations is not None:
                if has_annotations:
                    statement = (
                        statement
                        .join(AdminAnnotation, AdminAnnotation.document_id == AdminDocument.document_id)
                        .group_by(AdminDocument.document_id)
                    )
                else:
                    statement = (
                        statement
                        .outerjoin(AdminAnnotation, AdminAnnotation.document_id == AdminDocument.document_id)
                        .where(AdminAnnotation.annotation_id.is_(None))
                    )

            statement = statement.order_by(AdminDocument.created_at.desc())
            statement = statement.offset(offset).limit(limit)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results), total

    def update_status(
        self,
        document_id: str,
        status: str,
        auto_label_status: str | None = None,
        auto_label_error: str | None = None,
    ) -> None:
        """Update document status.

        Args:
            document_id: Document UUID as string
            status: New status
            auto_label_status: Auto-label status
            auto_label_error: Auto-label error message
        """
        with get_session_context() as session:
            document = session.get(AdminDocument, UUID(document_id))
            if document:
                document.status = status
                document.updated_at = datetime.now(timezone.utc)
                if auto_label_status is not None:
                    document.auto_label_status = auto_label_status
                if auto_label_error is not None:
                    document.auto_label_error = auto_label_error
                session.add(document)

    def update_file_path(self, document_id: str, file_path: str) -> None:
        """Update document file path."""
        with get_session_context() as session:
            document = session.get(AdminDocument, UUID(document_id))
            if document:
                document.file_path = file_path
                document.updated_at = datetime.now(timezone.utc)
                session.add(document)

    def update_group_key(self, document_id: str, group_key: str | None) -> bool:
        """Update document group key."""
        with get_session_context() as session:
            document = session.get(AdminDocument, UUID(document_id))
            if document:
                document.group_key = group_key
                document.updated_at = datetime.now(timezone.utc)
                session.add(document)
                return True
            return False

    def update_category(self, document_id: str, category: str) -> AdminDocument | None:
        """Update document category."""
        with get_session_context() as session:
            document = session.get(AdminDocument, UUID(document_id))
            if document:
                document.category = category
                document.updated_at = datetime.now(timezone.utc)
                session.add(document)
                session.commit()
                session.refresh(document)
                return document
            return None

    def delete(self, document_id: str) -> bool:
        """Delete a document and its annotations.

        Args:
            document_id: Document UUID as string

        Returns:
            True if deleted, False if not found
        """
        with get_session_context() as session:
            document = session.get(AdminDocument, UUID(document_id))
            if document:
                ann_stmt = select(AdminAnnotation).where(
                    AdminAnnotation.document_id == UUID(document_id)
                )
                annotations = session.exec(ann_stmt).all()
                for ann in annotations:
                    session.delete(ann)
                session.delete(document)
                session.commit()
                return True
            return False

    def get_categories(self) -> list[str]:
        """Get list of unique document categories."""
        with get_session_context() as session:
            statement = (
                select(AdminDocument.category)
                .distinct()
                .order_by(AdminDocument.category)
            )
            categories = session.exec(statement).all()
            return [c for c in categories if c is not None]

    def get_labeled_for_export(
        self,
        admin_token: str | None = None,
    ) -> list[AdminDocument]:
        """Get all labeled documents ready for export."""
        with get_session_context() as session:
            statement = select(AdminDocument).where(
                AdminDocument.status == "labeled"
            )
            if admin_token:
                statement = statement.where(AdminDocument.admin_token == admin_token)
            statement = statement.order_by(AdminDocument.created_at)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)
            return list(results)

    def count_by_status(
        self,
        admin_token: str | None = None,
    ) -> dict[str, int]:
        """Count documents by status."""
        with get_session_context() as session:
            statement = select(
                AdminDocument.status,
                func.count(AdminDocument.document_id),
            ).group_by(AdminDocument.status)

            results = session.exec(statement).all()
            return {status: count for status, count in results}

    def get_by_ids(self, document_ids: list[str]) -> list[AdminDocument]:
        """Get documents by list of IDs."""
        with get_session_context() as session:
            uuids = [UUID(str(did)) for did in document_ids]
            results = session.exec(
                select(AdminDocument).where(AdminDocument.document_id.in_(uuids))
            ).all()
            for r in results:
                session.expunge(r)
            return list(results)

    def get_for_training(
        self,
        admin_token: str | None = None,
        status: str = "labeled",
        has_annotations: bool = True,
        min_annotation_count: int | None = None,
        exclude_used_in_training: bool = False,
        limit: int = 100,
        offset: int = 0,
    ) -> tuple[list[AdminDocument], int]:
        """Get documents suitable for training with filtering."""
        from backend.data.admin_models import TrainingDocumentLink

        with get_session_context() as session:
            statement = select(AdminDocument).where(
                AdminDocument.status == status,
            )

            if has_annotations or min_annotation_count:
                annotation_subq = (
                    select(func.count(AdminAnnotation.annotation_id))
                    .where(AdminAnnotation.document_id == AdminDocument.document_id)
                    .correlate(AdminDocument)
                    .scalar_subquery()
                )

                if has_annotations:
                    statement = statement.where(annotation_subq > 0)

                if min_annotation_count:
                    statement = statement.where(annotation_subq >= min_annotation_count)

            if exclude_used_in_training:
                from sqlalchemy import exists
                training_subq = exists(
                    select(1)
                    .select_from(TrainingDocumentLink)
                    .where(TrainingDocumentLink.document_id == AdminDocument.document_id)
                )
                statement = statement.where(~training_subq)

            count_statement = select(func.count()).select_from(statement.subquery())
            total = session.exec(count_statement).one()

            statement = statement.order_by(AdminDocument.created_at.desc())
            statement = statement.limit(limit).offset(offset)

            results = session.exec(statement).all()
            for r in results:
                session.expunge(r)

            return list(results), total

    def acquire_annotation_lock(
        self,
        document_id: str,
        admin_token: str | None = None,
        duration_seconds: int = 300,
    ) -> AdminDocument | None:
        """Acquire annotation lock for a document."""
        from datetime import timedelta

        with get_session_context() as session:
            doc = session.get(AdminDocument, UUID(document_id))
            if not doc:
                return None

            now = datetime.now(timezone.utc)
            lock_until = doc.annotation_lock_until
            # Handle PostgreSQL returning offset-naive datetimes
            if lock_until and lock_until.tzinfo is None:
                lock_until = lock_until.replace(tzinfo=timezone.utc)
            if lock_until and lock_until > now:
                return None

            doc.annotation_lock_until = now + timedelta(seconds=duration_seconds)
            session.add(doc)
            session.commit()
            session.refresh(doc)
            session.expunge(doc)
            return doc

    def release_annotation_lock(
        self,
        document_id: str,
        admin_token: str | None = None,
        force: bool = False,
    ) -> AdminDocument | None:
        """Release annotation lock for a document."""
        with get_session_context() as session:
            doc = session.get(AdminDocument, UUID(document_id))
            if not doc:
                return None

            doc.annotation_lock_until = None
            session.add(doc)
            session.commit()
            session.refresh(doc)
            session.expunge(doc)
            return doc

    def extend_annotation_lock(
        self,
        document_id: str,
        admin_token: str | None = None,
        additional_seconds: int = 300,
    ) -> AdminDocument | None:
        """Extend an existing annotation lock."""
        from datetime import timedelta

        with get_session_context() as session:
            doc = session.get(AdminDocument, UUID(document_id))
            if not doc:
                return None

            now = datetime.now(timezone.utc)
            lock_until = doc.annotation_lock_until
            # Handle PostgreSQL returning offset-naive datetimes
            if lock_until and lock_until.tzinfo is None:
                lock_until = lock_until.replace(tzinfo=timezone.utc)
            if not lock_until or lock_until <= now:
                return None

            doc.annotation_lock_until = lock_until + timedelta(seconds=additional_seconds)
            session.add(doc)
            session.commit()
            session.refresh(doc)
            session.expunge(doc)
            return doc
