OCR/docengine/app/repositories/document_repository.py

from __future__ import annotations

import uuid

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.document import (
    Document,
    DocumentImage,
    DocumentPage,
    DocumentTable,
    DocumentTextBlock,
    TemplateMatch,
)
from app.repositories.base import BaseRepository


class DocumentRepository(BaseRepository[Document]):
    """Repository for Document operations."""

    def __init__(self, db: Session) -> None:
        super().__init__(db, Document)

    def get_by_checksum(self, checksum: str) -> Document | None:
        """Get document by file checksum."""
        query = select(Document).where(Document.checksum == checksum)
        result = self.db.execute(query)
        return result.scalars().first()

    def get_by_status(self, status: str, offset: int = 0, limit: int = 100) -> list[Document]:
        """Get documents by processing status."""
        query = (
            select(Document)
            .where(Document.status == status)
            .order_by(Document.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        result = self.db.execute(query)
        return list(result.scalars().all())

    def get_user_documents(
        self,
        user_id: uuid.UUID,
        offset: int = 0,
        limit: int = 100,
    ) -> list[Document]:
        """Get documents uploaded by a specific user."""
        query = (
            select(Document)
            .where(Document.uploaded_by == user_id)
            .order_by(Document.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        result = self.db.execute(query)
        return list(result.scalars().all())

    def update_status(
        self,
        document_id: uuid.UUID,
        status: str,
        error_message: str | None = None,
    ) -> Document | None:
        """Update document processing status."""
        document = self.get_by_id(document_id)
        if document:
            document.status = status
            if error_message:
                document.error_message = error_message
            self.db.flush()
            self.db.refresh(document)
        return document

    def get_with_pages(self, document_id: uuid.UUID) -> Document | None:
        """Get document with all pages eagerly loaded."""
        return self.get_by_id(document_id)

    def get_pending_documents(self, limit: int = 10) -> list[Document]:
        """Get pending documents for processing."""
        return self.get_by_status("pending", limit=limit)


class DocumentPageRepository(BaseRepository[DocumentPage]):
    """Repository for DocumentPage operations."""

    def __init__(self, db: Session) -> None:
        super().__init__(db, DocumentPage)

    def get_document_pages(self, document_id: uuid.UUID) -> list[DocumentPage]:
        """Get all pages for a document ordered by page number."""
        query = (
            select(DocumentPage)
            .where(DocumentPage.document_id == document_id)
            .order_by(DocumentPage.page_number)
        )
        result = self.db.execute(query)
        return list(result.scalars().all())

    def get_page_by_number(self, document_id: uuid.UUID, page_number: int) -> DocumentPage | None:
        """Get a specific page by document ID and page number."""
        query = select(DocumentPage).where(
            DocumentPage.document_id == document_id,
            DocumentPage.page_number == page_number,
        )
        result = self.db.execute(query)
        return result.scalars().first()

    def create_page(
        self,
        document_id: uuid.UUID,
        page_number: int,
        width: float,
        height: float,
        image_path: str | None = None,
        text_content: str | None = None,
    ) -> DocumentPage:
        """Create a new document page."""
        page = DocumentPage(
            document_id=document_id,
            page_number=page_number,
            width=width,
            height=height,
            image_path=image_path,
            text_content=text_content,
        )
        return self.create(page)


class DocumentTextBlockRepository(BaseRepository[DocumentTextBlock]):
    """Repository for DocumentTextBlock operations."""

    def __init__(self, db: Session) -> None:
        super().__init__(db, DocumentTextBlock)

    def get_page_text_blocks(self, page_id: uuid.UUID) -> list[DocumentTextBlock]:
        """Get all text blocks for a page."""
        query = (
            select(DocumentTextBlock)
            .where(DocumentTextBlock.page_id == page_id)
            .order_by(DocumentTextBlock.sequence)
        )
        result = self.db.execute(query)
        return list(result.scalars().all())

    def get_by_block_type(self, page_id: uuid.UUID, block_type: str) -> list[DocumentTextBlock]:
        """Get text blocks by type (header, footer, watermark, text)."""
        query = (
            select(DocumentTextBlock)
            .where(
                DocumentTextBlock.page_id == page_id,
                DocumentTextBlock.block_type == block_type,
            )
            .order_by(DocumentTextBlock.sequence)
        )
        result = self.db.execute(query)
        return list(result.scalars().all())

    def create_text_block(
        self,
        page_id: uuid.UUID,
        text: str,
        x: float,
        y: float,
        width: float,
        height: float,
        confidence: float | None = None,
        font_family: str | None = None,
        font_size: float | None = None,
        font_color: str | None = None,
        font_style: str | None = None,
        block_type: str = "text",
        sequence: int = 0,
    ) -> DocumentTextBlock:
        """Create a new text block."""
        text_block = DocumentTextBlock(
            page_id=page_id,
            text=text,
            x=x,
            y=y,
            width=width,
            height=height,
            confidence=confidence,
            font_family=font_family,
            font_size=font_size,
            font_color=font_color,
            font_style=font_style,
            block_type=block_type,
            sequence=sequence,
        )
        return self.create(text_block)


class DocumentImageRepository(BaseRepository[DocumentImage]):
    """Repository for DocumentImage operations."""

    def __init__(self, db: Session) -> None:
        super().__init__(db, DocumentImage)

    def get_page_images(self, page_id: uuid.UUID) -> list[DocumentImage]:
        """Get all images for a page."""
        query = select(DocumentImage).where(DocumentImage.page_id == page_id)
        result = self.db.execute(query)
        return list(result.scalars().all())

    def create_image(
        self,
        page_id: uuid.UUID,
        x: float,
        y: float,
        width: float,
        height: float,
        image_path: str,
        image_type: str = "figure",
    ) -> DocumentImage:
        """Create a new document image record."""
        image = DocumentImage(
            page_id=page_id,
            x=x,
            y=y,
            width=width,
            height=height,
            image_path=image_path,
            image_type=image_type,
        )
        return self.create(image)


class DocumentTableRepository(BaseRepository[DocumentTable]):
    """Repository for DocumentTable operations."""

    def __init__(self, db: Session) -> None:
        super().__init__(db, DocumentTable)

    def get_page_tables(self, page_id: uuid.UUID) -> list[DocumentTable]:
        """Get all tables for a page."""
        query = select(DocumentTable).where(DocumentTable.page_id == page_id)
        result = self.db.execute(query)
        return list(result.scalars().all())

    def create_table(
        self,
        page_id: uuid.UUID,
        x: float,
        y: float,
        width: float,
        height: float,
        rows: int,
        columns: int,
        data: dict | None = None,
    ) -> DocumentTable:
        """Create a new document table record."""
        table = DocumentTable(
            page_id=page_id,
            x=x,
            y=y,
            width=width,
            height=height,
            rows=rows,
            columns=columns,
            data=data,
        )
        return self.create(table)


class TemplateMatchRepository(BaseRepository[TemplateMatch]):
    """Repository for TemplateMatch operations."""

    def __init__(self, db: Session) -> None:
        super().__init__(db, TemplateMatch)

    def get_document_matches(
        self,
        document_id: uuid.UUID,
        min_confidence: float = 0.0,
    ) -> list[TemplateMatch]:
        """Get all template matches for a document."""
        query = (
            select(TemplateMatch)
            .where(
                TemplateMatch.document_id == document_id,
                TemplateMatch.confidence_score >= min_confidence,
            )
            .order_by(TemplateMatch.confidence_score.desc())
        )
        result = self.db.execute(query)
        return list(result.scalars().all())

    def get_selected_match(self, document_id: uuid.UUID) -> TemplateMatch | None:
        """Get the selected template match for a document."""
        query = select(TemplateMatch).where(
            TemplateMatch.document_id == document_id,
            TemplateMatch.selected.is_(True),
        )
        result = self.db.execute(query)
        return result.scalars().first()

    def select_match(self, match_id: uuid.UUID) -> TemplateMatch | None:
        """Select a template match (deselecting all others for the same document)."""
        match = self.get_by_id(match_id)
        if not match:
            return None

        # Deselect all other matches for this document
        query = select(TemplateMatch).where(
            TemplateMatch.document_id == match.document_id,
            TemplateMatch.selected.is_(True),
        )
        result = self.db.execute(query)
        for existing_match in result.scalars().all():
            existing_match.selected = False

        match.selected = True
        self.db.flush()
        self.db.refresh(match)
        return match

    def create_match(
        self,
        document_id: uuid.UUID,
        format_id: uuid.UUID,
        confidence_score: float,
        match_details: dict | None = None,
        selected: bool = False,
    ) -> TemplateMatch:
        """Create a new template match."""
        template_match = TemplateMatch(
            document_id=document_id,
            format_id=format_id,
            confidence_score=confidence_score,
            match_details=match_details,
            selected=selected,
        )
        return self.create(template_match)