from __future__ import annotations
import asyncpg
from ..models import ExtractionResult, Fact, Citation


def _row_to_fact(f: dict) -> Fact:
    return Fact(
        id=f["id"],
        text=f["text"],
        citations=[Citation(**c) for c in f["citations"]],
    )


async def insert_extraction_result(
    pool: asyncpg.Pool,
    document_id: str,
    facts: list[Fact],
    prompt: str | None = None,
    no_mention_found: bool = False,
    absence_evidence: str | None = None,
    source: str = "extraction",
) -> ExtractionResult:
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            INSERT INTO extraction_results
                (document_id, facts, prompt, no_mention_found, absence_evidence, source)
            VALUES ($1,$2,$3,$4,$5,$6)
            RETURNING id, extracted_at
            """,
            document_id,
            [f.model_dump() for f in facts],
            prompt,
            no_mention_found,
            absence_evidence,
            source,
        )

    return ExtractionResult(
        id=str(row["id"]),
        documentId=document_id,
        facts=facts,
        prompt=prompt,
        noMentionFound=no_mention_found,
        absenceEvidence=absence_evidence,
        extractedAt=row["extracted_at"].isoformat(),
    )


async def get_latest_extraction_by_document_id(
    pool: asyncpg.Pool, document_id: str
) -> ExtractionResult | None:
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            SELECT id, document_id, facts, prompt, no_mention_found, absence_evidence, extracted_at
            FROM extraction_results
            WHERE document_id = $1
            ORDER BY extracted_at DESC LIMIT 1
            """,
            document_id,
        )

    if row is None:
        return None

    facts = [_row_to_fact(f) for f in (row["facts"] or [])]
    return ExtractionResult(
        id=str(row["id"]),
        documentId=row["document_id"],
        facts=facts,
        prompt=row["prompt"],
        noMentionFound=row["no_mention_found"],
        absenceEvidence=row["absence_evidence"],
        extractedAt=row["extracted_at"].isoformat(),
    )


async def get_all_facts_by_document_id(
    pool: asyncpg.Pool, document_id: str
) -> ExtractionResult | None:
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            """
            SELECT facts, extracted_at
            FROM extraction_results
            WHERE document_id = $1
            ORDER BY extracted_at ASC
            """,
            document_id,
        )

    if not rows:
        return None

    all_facts: list[Fact] = [
        _row_to_fact(f) for row in rows for f in (row["facts"] or [])
    ]
    latest_at = rows[-1]["extracted_at"]

    return ExtractionResult(
        documentId=document_id,
        facts=all_facts,
        extractedAt=latest_at.isoformat(),
    )
