"""Page-level evidence typing, target-attributed citations and review state.

Revision ID: 0006
Revises: 0005

Existing indexes are deliberately not upgraded in place. Chunking changed (a contents
page is now stored as navigation rather than body text), so legacy documents keep their
old ``index_version`` and are excluded from retrieval until the reindex task rebuilds
them. That is what stops an old index from being silently used as classified body text.
Reclassification of page types needs no embedding model, so this migration applies it
immediately for every document whose parsed pages are still saved on disk.
"""
import json

from alembic import op
import sqlalchemy as sa

revision = "0006"
down_revision = "0005"
branch_labels = None
depends_on = None


PAGE_COLUMNS = (
    ("page_type", sa.String(20), "unknown"),
    ("mixed", sa.Boolean(), "false"),
)
PAGE_TEXT_COLUMNS = (
    ("section", sa.String(500), ""),
    ("chapter", sa.String(500), ""),
)
CHUNK_COLUMNS = (
    ("chapter", sa.String(500), ""),
    ("page_type", sa.String(20), "unknown"),
    ("block_type", sa.String(20), "legacy"),
)
EVIDENCE_COLUMNS = (
    ("target_id", sa.String(50), ""),
    ("target_label", sa.String(300), ""),
    ("page_type", sa.String(20), "unknown"),
    ("source_role", sa.String(20), "content"),
    ("historical", sa.Boolean(), "false"),
)


def _columns(table):
    return {column["name"] for column in sa.inspect(op.get_bind()).get_columns(table)}


def _add(table, name, column_type, default, *, nullable=False):
    if name in _columns(table):
        return False
    op.add_column(table, sa.Column(name, column_type, nullable=nullable, server_default=default))
    return True


def upgrade():
    added = {}
    for name, column_type, default in PAGE_COLUMNS:
        added[f"knowledge_pages.{name}"] = _add("knowledge_pages", name, column_type, default)
    added["knowledge_pages.section"] = _add("knowledge_pages", "section", sa.String(500), "")
    added["knowledge_pages.chapter"] = _add("knowledge_pages", "chapter", sa.String(500), "")
    added["knowledge_pages.confidence"] = _add("knowledge_pages", "confidence", sa.Float(), None, nullable=True)
    added["knowledge_pages.classification"] = _add("knowledge_pages", "classification", sa.JSON(), "{}")
    for name, column_type, default in CHUNK_COLUMNS:
        added[f"knowledge_chunks.{name}"] = _add("knowledge_chunks", name, column_type, default)
    for name, column_type, default in EVIDENCE_COLUMNS:
        added[f"score_evidences.{name}"] = _add("score_evidences", name, column_type, default)
    added["criterion_scores.evidence_gaps"] = _add("criterion_scores", "evidence_gaps", sa.JSON(), "[]")

    db = op.get_bind()
    # Quotations written before this policy existed are marked as historical so the
    # interface never presents a contents-page quotation as report body evidence.
    db.execute(sa.text("UPDATE score_evidences SET historical = true WHERE historical = false"))
    # A legacy chunk was never split by block, so it may not be used as content
    # evidence until the reindex task rebuilds it.
    if added["knowledge_chunks.block_type"]:
        db.execute(sa.text("UPDATE knowledge_chunks SET block_type = 'legacy'"))
    _reclassify_saved_pages(db)


def _json(value):
    return json.loads(value) if isinstance(value, str) else value or {}


def _reclassify_saved_pages(db):
    """Re-derive page types from saved parse output; retryable and audit friendly."""
    from wisegrade.pages import classify_document

    files = {row.id: _json(row.parsed) for row in db.execute(sa.text(
        "SELECT id, parsed FROM submission_files")).mappings()}
    documents = db.execute(sa.text(
        "SELECT id, submission_file_id FROM knowledge_documents WHERE submission_file_id IS NOT NULL")).mappings().all()
    update = sa.text("UPDATE knowledge_pages SET page_type = :page_type, mixed = :mixed, section = :section, "
                     "chapter = :chapter, confidence = :confidence, classification = :classification "
                     "WHERE document_id = :document_id AND page = :page")
    chunk_update = sa.text("UPDATE knowledge_chunks SET page_type = :page_type, section = :section, "
                           "chapter = :chapter WHERE document_id = :document_id AND page = :page")
    rebuilt = 0
    for doc in documents:
        pages = files.get(doc.submission_file_id, {}).get("pages", [])
        if not pages:
            continue
        for item in classify_document(pages):
            payload = {"page_type": item["page_type"], "mixed": bool(item["mixed"]),
                       "section": item["section"], "chapter": item["chapter"],
                       "confidence": item["confidence"],
                       "classification": json.dumps({**item["signals"],
                                                     "classifier_version": item["classifier_version"]},
                                                    ensure_ascii=False, default=str),
                       "document_id": doc.id, "page": item["page"]}
            db.execute(update, payload)
            db.execute(chunk_update, payload)
        rebuilt += 1
    if rebuilt:
        db.execute(sa.text("INSERT INTO audit_logs (action, entity_type, entity_id, details, created_at) "
                           "VALUES ('reclassify', 'knowledge_document', NULL, :details, CURRENT_TIMESTAMP)"),
                   {"details": json.dumps({"migration": "0006", "documents": rebuilt}, ensure_ascii=False)})


def downgrade():
    for table, names in (("criterion_scores", ("evidence_gaps",)),
                         ("score_evidences", tuple(name for name, _, _ in EVIDENCE_COLUMNS)),
                         ("knowledge_chunks", tuple(name for name, _, _ in CHUNK_COLUMNS)),
                         ("knowledge_pages", ("page_type", "mixed", "section", "chapter", "confidence", "classification"))):
        for name in names:
            if name in _columns(table):
                op.drop_column(table, name)