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)