"""Memory single-chain cleanup and memoir state control columns.

Revision ID: 0015_memory_single_chain
Revises: 0014_memory_evidence_indexes
"""

from typing import Sequence, Union

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

revision: str = "0015_memory_single_chain"
down_revision: Union[str, None] = "0014_memory_evidence_indexes"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _has_column(table: str, column: str) -> bool:
    bind = op.get_bind()
    return any(c["name"] == column for c in sa.inspect(bind).get_columns(table))


def upgrade() -> None:
    if not _has_column("memoir_states", "known_facts_json"):
        op.add_column(
            "memoir_states",
            sa.Column(
                "known_facts_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )
    if not _has_column("memoir_states", "persona_threads_json"):
        op.add_column(
            "memoir_states",
            sa.Column(
                "persona_threads_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )
    if not _has_column("memoir_states", "recent_questions_json"):
        op.add_column(
            "memoir_states",
            sa.Column(
                "recent_questions_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )

    op.execute(
        """
        UPDATE memoir_states
        SET
          known_facts_json = COALESCE((slots::jsonb -> '__interview_state__' -> 'known_facts')::json, '[]'::json),
          persona_threads_json = COALESCE((slots::jsonb -> '__interview_state__' -> 'persona_threads')::json, '[]'::json),
          recent_questions_json = COALESCE((slots::jsonb -> '__interview_state__' -> 'recent_questions')::json, '[]'::json),
          slots = (slots::jsonb - '__interview_state__')::json
        WHERE slots::jsonb ? '__interview_state__'
        """
    )

    if _has_column("chapters", "evidence_bundle_json"):
        op.drop_column("chapters", "evidence_bundle_json")


def downgrade() -> None:
    if not _has_column("chapters", "evidence_bundle_json"):
        op.add_column(
            "chapters",
            sa.Column(
                "evidence_bundle_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )
    for column in (
        "recent_questions_json",
        "persona_threads_json",
        "known_facts_json",
    ):
        if _has_column("memoir_states", column):
            op.drop_column("memoir_states", column)