from __future__ import annotations
import asyncpg
from ..models import Document, DocumentSummary, TextSegment, DocumentContent


async def insert_document(pool: asyncpg.Pool, doc: Document) -> None:
    async with pool.acquire() as conn:
        await conn.execute(
            """
            INSERT INTO documents (
                id, title, type, raw_text, segments, page_count, page_dimensions,
                source_url, original_file_name, pdf_data, created_at
            ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)
            """,
            doc.id,
            doc.title,
            doc.type,
            doc.content.rawText,
            [s.model_dump() for s in doc.content.segments],
            doc.content.pageCount,
            doc.content.pageDimensions,
            doc.sourceUrl,
            doc.originalFileName,
            doc.pdfData,
            doc.createdAt,
        )


async def get_document_by_id(pool: asyncpg.Pool, doc_id: str) -> Document | None:
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            SELECT id, title, type, raw_text, segments, page_count, page_dimensions,
                   source_url, original_file_name, pdf_data, created_at
            FROM documents WHERE id = $1
            """,
            doc_id,
        )

    if row is None:
        return None

    segments = [TextSegment(**s) for s in (row["segments"] or [])]

    return Document(
        id=row["id"],
        title=row["title"],
        type=row["type"],
        content=DocumentContent(
            rawText=row["raw_text"],
            segments=segments,
            pageCount=row["page_count"],
            pageDimensions=row["page_dimensions"],
        ),
        sourceUrl=row["source_url"],
        originalFileName=row["original_file_name"],
        pdfData=row["pdf_data"],
        createdAt=row["created_at"].isoformat(),
    )


async def list_documents(pool: asyncpg.Pool) -> list[DocumentSummary]:
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            "SELECT id, title, type, source_url, created_at FROM documents ORDER BY created_at DESC"
        )

    return [
        DocumentSummary(
            id=r["id"],
            title=r["title"],
            type=r["type"],
            sourceUrl=r["source_url"],
            createdAt=r["created_at"].isoformat(),
        )
        for r in rows
    ]
