import asyncpg


async def init_db(pool: asyncpg.Pool) -> None:
    async with pool.acquire() as conn:
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id                 TEXT PRIMARY KEY,
                title              TEXT NOT NULL,
                type               TEXT NOT NULL,
                raw_text           TEXT NOT NULL,
                segments           JSONB NOT NULL DEFAULT '[]',
                page_count         INTEGER,
                page_dimensions    JSONB,
                source_url         TEXT,
                original_file_name TEXT,
                pdf_data           TEXT,
                created_at         TIMESTAMPTZ NOT NULL DEFAULT NOW()
            )
        """)

        await conn.execute("""
            CREATE TABLE IF NOT EXISTS extraction_results (
                id           UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                document_id  TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
                facts        JSONB NOT NULL DEFAULT '[]',
                prompt       TEXT,
                extracted_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
            )
        """)

        await conn.execute("""
            CREATE INDEX IF NOT EXISTS extraction_results_document_id_idx
                ON extraction_results(document_id)
        """)

        for stmt in [
            "ALTER TABLE extraction_results ADD COLUMN IF NOT EXISTS no_mention_found BOOLEAN NOT NULL DEFAULT FALSE",
            "ALTER TABLE extraction_results ADD COLUMN IF NOT EXISTS absence_evidence TEXT",
            "ALTER TABLE extraction_results ADD COLUMN IF NOT EXISTS source VARCHAR(20) NOT NULL DEFAULT 'extraction'",
        ]:
            await conn.execute(stmt)
