src/utils/files.py

import datetime
import logging
from io import BytesIO
from typing import Any, Protocol

from fastapi import UploadFile
from nanoid import generate as generate_nanoid
from sqlalchemy import Integer, select
from sqlalchemy.ext.asyncio import AsyncSession

from src import schemas
from src.config import settings
from src.exceptions import (
    FileProcessingError,
    UnsupportedFileTypeError,
    ValidationException,
)
from src.schemas import Message

logger = logging.getLogger(__name__)


class FileProcessor(Protocol):
    async def extract_text(self, content: bytes) -> str: ...
    def supports_file_type(self, content_type: str) -> bool: ...


class PDFProcessor:
    def supports_file_type(self, content_type: str) -> bool:
        return content_type == "application/pdf"

    async def extract_text(self, content: bytes) -> str:
        import pdfplumber

        with pdfplumber.open(BytesIO(content)) as pdf_reader:
            text_parts: list[str] = []
            for page_num, page in enumerate(pdf_reader.pages):
                text = page.extract_text()
                if text and text.strip():
                    text_parts.append(f"[Page {page_num + 1}]\n{text}")
            return "\n\n".join(text_parts)


class TextProcessor:
    def supports_file_type(self, content_type: str) -> bool:
        return content_type.startswith("text/")

    async def extract_text(self, content: bytes) -> str:
        # Try different encodings
        for encoding in ["utf-8", "utf-16", "latin-1"]:
            try:
                return content.decode(encoding)
            except UnicodeDecodeError:
                continue
        raise ValueError("Could not decode text file")


class JSONProcessor:
    def supports_file_type(self, content_type: str) -> bool:
        return content_type == "application/json"

    async def extract_text(self, content: bytes) -> str:
        import json

        try:
            decoded_content = content.decode("utf-8")
        except UnicodeDecodeError as exc:
            raise ValidationException("JSON uploads must be UTF-8 encoded") from exc

        if not decoded_content.strip():
            return ""

        try:
            data = json.loads(decoded_content)
        except json.JSONDecodeError as exc:
            raise ValidationException("Uploaded JSON is invalid") from exc

        # Convert JSON to readable text format
        return json.dumps(data, ensure_ascii=False)


class FileProcessingService:
    def __init__(self):
        self.processors: list[FileProcessor] = [
            PDFProcessor(),
            TextProcessor(),
            JSONProcessor(),
            # Add more processors as needed
        ]

    async def extract_text_from_upload(self, file: UploadFile) -> str:
        """Extract text from uploaded file without saving to disk."""
        content = await file.read()

        # Reset file position in case it's needed again
        await file.seek(0)

        processor = self._get_processor(file.content_type or "")
        if not processor:
            raise UnsupportedFileTypeError(
                f"Unsupported file type: {file.content_type}. Supported types: {[p.__class__.__name__ for p in self.processors]}"
            )

        return await processor.extract_text(content)

    def _get_processor(self, content_type: str) -> FileProcessor | None:
        for processor in self.processors:
            if processor.supports_file_type(content_type):
                return processor
        return None


def split_text_into_chunks(text: str, max_chars: int = 49500) -> list[str]:
    """Split text into chunks that fit within message limits."""
    if len(text) <= max_chars:
        return [text]

    chunks: list[str] = []
    current_pos = 0

    while current_pos < len(text):
        # Try to break at paragraph boundaries first
        end_pos = current_pos + max_chars

        if end_pos >= len(text):
            chunks.append(text[current_pos:])
            break

        # Look for good break points (paragraph, sentence, word)
        break_pos = end_pos
        for delimiter in ["\n\n", "\n", ". ", " "]:
            last_delimiter = text.rfind(delimiter, current_pos, end_pos)
            if last_delimiter > current_pos:
                break_pos = last_delimiter + len(delimiter)
                break

        chunks.append(text[current_pos:break_pos])
        current_pos = break_pos

    return chunks


async def get_file_messages(
    db: AsyncSession,
    workspace_name: str,
    file_id: str,
    session_name: str | None = None,
) -> list[Message]:
    """Get all messages for a specific document, ordered by chunk_index."""
    from sqlalchemy import and_, func

    from src.models import Message

    query = select(Message).where(
        and_(
            Message.workspace_name == workspace_name,
            func.jsonb_extract_path_text(Message.internal_metadata, "file_id")
            == file_id,
        )
    )

    if session_name:
        query = query.where(Message.session_name == session_name)

    # Order by chunk_index
    query = query.order_by(
        func.jsonb_extract_path_text(Message.internal_metadata, "chunk_index").cast(
            Integer
        )
    )

    result = await db.execute(query)
    return list(result.scalars().all())


async def process_file_uploads_for_messages(
    file: UploadFile,
    peer_id: str,
    max_chars: int = settings.MAX_MESSAGE_SIZE,
    metadata: dict[str, Any] | None = None,
    configuration: schemas.MessageConfiguration | None = None,
    created_at: datetime.datetime | None = None,
) -> list[dict[str, Any]]:
    """
    Process an uploaded file and prepare message creation data.

    This function extracts text from a file, splits it into chunks, and prepares
    the data needed to create messages.

    Args:
        file: Uploaded file to process
        peer_id: ID of the peer creating the messages
        max_chars: Maximum characters per message chunk
        metadata: Optional metadata to associate with all messages created from this file
        configuration: Optional configuration to associate with all messages created from this file
        created_at: Optional created_at timestamp to use for all messages created from this file

    Returns:
        List of dictionaries containing message_create and file_metadata

    Raises:
        HTTPException: If file processing fails
    """

    file_processor = FileProcessingService()
    all_message_data: list[dict[str, Any]] = []

    # Process the uploaded file
    extracted_text = await file_processor.extract_text_from_upload(file)

    # Split into chunks and create messages
    chunks = split_text_into_chunks(extracted_text, max_chars=max_chars)
    file_id = generate_nanoid()

    for i, chunk in enumerate(chunks):
        # Build message content properly handling empty files
        message_content = chunk or ""

        # Create message with optional metadata, configuration, and created_at
        message_create = schemas.MessageCreate(
            content=message_content,
            peer_id=peer_id,
            metadata=metadata,
            configuration=configuration,
            created_at=created_at,
        )

        # Store file metadata separately to add to internal_metadata later
        file_metadata = {
            "file_id": file_id,
            "filename": file.filename,
            "chunk_index": i,
            "total_chunks": len(chunks),
            "original_file_size": file.size,
            "content_type": file.content_type,
            "chunk_character_range": [
                i * max_chars,
                min((i + 1) * max_chars, len(extracted_text)),
            ],
        }

        all_message_data.append(
            {
                "message_create": message_create,
                "file_metadata": file_metadata,
            }
        )

    if not all_message_data:
        raise FileProcessingError()

    return all_message_data