"""Add conversational branching tables and columns.

Revision ID: 067
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID, JSONB

revision = "067"
down_revision = "066"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # session_branches
    op.create_table(
        "session_branches",
        sa.Column("id", UUID(as_uuid=True), primary_key=True),
        sa.Column("session_id", UUID(as_uuid=True), sa.ForeignKey("ai_sessions.id", ondelete="CASCADE"), nullable=False),
        sa.Column("parent_branch_id", UUID(as_uuid=True), sa.ForeignKey("session_branches.id", ondelete="CASCADE"), nullable=True),
        sa.Column("fork_point_step_id", UUID(as_uuid=True), sa.ForeignKey("ai_session_steps.id", ondelete="SET NULL"), nullable=True),
        sa.Column("branch_order", sa.Integer, nullable=False, server_default="1"),
        sa.Column("label", sa.String(200), nullable=False),
        sa.Column("status", sa.String(20), nullable=False, server_default="active"),
        sa.Column("status_reason", sa.Text, nullable=True),
        sa.Column("status_changed_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("status_changed_by", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=True),
        sa.Column("conversation_messages", JSONB, nullable=False, server_default="[]"),
        sa.Column("context_summary", JSONB, nullable=True),
        sa.Column("evidence_from_branch_id", UUID(as_uuid=True), sa.ForeignKey("session_branches.id", ondelete="SET NULL"), nullable=True),
        sa.Column("evidence_description", sa.Text, nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.CheckConstraint("status IN ('active', 'dead_end', 'solved', 'untried', 'revived')", name="ck_session_branches_status"),
        sa.CheckConstraint("branch_order > 0", name="ck_session_branches_branch_order_positive"),
    )
    op.create_index("ix_session_branches_session_id", "session_branches", ["session_id"])
    op.create_index("ix_session_branches_parent_branch_id", "session_branches", ["parent_branch_id"])
    op.create_index("ix_session_branches_session_status", "session_branches", ["session_id", "status"])
    op.create_index("ix_session_branches_session_order", "session_branches", ["session_id", "branch_order"])

    # fork_points
    op.create_table(
        "fork_points",
        sa.Column("id", UUID(as_uuid=True), primary_key=True),
        sa.Column("session_id", UUID(as_uuid=True), sa.ForeignKey("ai_sessions.id", ondelete="CASCADE"), nullable=False),
        sa.Column("parent_branch_id", UUID(as_uuid=True), sa.ForeignKey("session_branches.id", ondelete="CASCADE"), nullable=False),
        sa.Column("trigger_step_id", UUID(as_uuid=True), sa.ForeignKey("ai_session_steps.id", ondelete="SET NULL"), nullable=True),
        sa.Column("fork_reason", sa.Text, nullable=False),
        sa.Column("options", JSONB, nullable=False, server_default="[]"),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
    )
    op.create_index("ix_fork_points_session_id", "fork_points", ["session_id"])

    # session_handoffs
    op.create_table(
        "session_handoffs",
        sa.Column("id", UUID(as_uuid=True), primary_key=True),
        sa.Column("session_id", UUID(as_uuid=True), sa.ForeignKey("ai_sessions.id", ondelete="CASCADE"), nullable=False),
        sa.Column("handed_off_by", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
        sa.Column("intent", sa.String(20), nullable=False),
        sa.Column("source_branch_id", UUID(as_uuid=True), sa.ForeignKey("session_branches.id", ondelete="SET NULL"), nullable=True),
        sa.Column("snapshot", JSONB, nullable=False, server_default="{}"),
        sa.Column("ai_assessment", sa.Text, nullable=True),
        sa.Column("ai_assessment_data", JSONB, nullable=True),
        sa.Column("artifacts", JSONB, nullable=True),
        sa.Column("engineer_notes", sa.Text, nullable=True),
        sa.Column("priority", sa.String(20), nullable=False, server_default="normal"),
        sa.Column("claimed_by", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=True),
        sa.Column("claimed_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("psa_note_pushed", sa.Boolean, server_default="false"),
        sa.Column("psa_note_id", sa.String(100), nullable=True),
        sa.Column("notification_sent", sa.Boolean, server_default="false"),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.CheckConstraint("intent IN ('park', 'escalate')", name="ck_session_handoffs_intent"),
        sa.CheckConstraint("priority IN ('normal', 'elevated')", name="ck_session_handoffs_priority"),
    )
    op.create_index("ix_session_handoffs_session_id", "session_handoffs", ["session_id"])

    # session_resolution_outputs
    op.create_table(
        "session_resolution_outputs",
        sa.Column("id", UUID(as_uuid=True), primary_key=True),
        sa.Column("session_id", UUID(as_uuid=True), sa.ForeignKey("ai_sessions.id", ondelete="CASCADE"), nullable=False),
        sa.Column("output_type", sa.String(30), nullable=False),
        sa.Column("generated_content", sa.Text, nullable=False),
        sa.Column("structured_data", JSONB, nullable=True),
        sa.Column("edited_content", sa.Text, nullable=True),
        sa.Column("status", sa.String(20), nullable=False, server_default="draft"),
        sa.Column("pushed_to", sa.String(50), nullable=True),
        sa.Column("pushed_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("pushed_reference", sa.String(200), nullable=True),
        sa.Column("generated_by_model", sa.String(50), nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.CheckConstraint("output_type IN ('psa_ticket_notes', 'knowledge_base', 'client_summary')", name="ck_session_resolution_outputs_output_type"),
        sa.CheckConstraint("status IN ('draft', 'approved', 'pushed', 'rejected')", name="ck_session_resolution_outputs_status"),
        sa.UniqueConstraint("session_id", "output_type", name="uq_session_resolution_session_type"),
    )
    op.create_index("ix_session_resolution_outputs_session_id", "session_resolution_outputs", ["session_id"])

    # ai_sessions: add 5 columns (NO FK on active_branch_id)
    op.add_column("ai_sessions", sa.Column("is_branching", sa.Boolean, server_default="false", nullable=False))
    op.add_column("ai_sessions", sa.Column("active_branch_id", UUID(as_uuid=True), nullable=True))
    op.add_column("ai_sessions", sa.Column("handoff_count", sa.Integer, server_default="0", nullable=False))
    op.add_column("ai_sessions", sa.Column("total_active_seconds", sa.Integer, server_default="0", nullable=False))
    op.add_column("ai_sessions", sa.Column("total_parked_seconds", sa.Integer, server_default="0", nullable=False))

    # ai_session_steps: add 3 columns + update CHECK
    op.add_column("ai_session_steps", sa.Column("branch_id", UUID(as_uuid=True), sa.ForeignKey("session_branches.id", ondelete="SET NULL"), nullable=True))
    op.add_column("ai_session_steps", sa.Column("is_fork_point", sa.Boolean, server_default="false", nullable=False))
    op.add_column("ai_session_steps", sa.Column("fork_point_id", UUID(as_uuid=True), sa.ForeignKey("fork_points.id", ondelete="SET NULL"), nullable=True))
    op.create_index("ix_ai_session_steps_branch_id", "ai_session_steps", ["branch_id"])

    op.drop_constraint("ck_ai_session_steps_step_type", "ai_session_steps", type_="check")
    op.create_check_constraint(
        "ck_ai_session_steps_step_type", "ai_session_steps",
        "step_type IN ('question', 'action', 'script_generation', 'verification', 'info_request', 'note', 'intake_analysis', 'fork')",
    )

    # file_uploads: add 5 columns
    op.add_column("file_uploads", sa.Column("ai_description", sa.Text, nullable=True))
    op.add_column("file_uploads", sa.Column("extracted_content", sa.Text, nullable=True))
    op.add_column("file_uploads", sa.Column("content_summary", sa.Text, nullable=True))
    op.add_column("file_uploads", sa.Column("uploaded_on_branch_id", UUID(as_uuid=True), sa.ForeignKey("session_branches.id", ondelete="SET NULL"), nullable=True))
    op.add_column("file_uploads", sa.Column("uploaded_at_step_id", UUID(as_uuid=True), sa.ForeignKey("ai_session_steps.id", ondelete="SET NULL"), nullable=True))


def downgrade() -> None:
    op.drop_column("file_uploads", "uploaded_at_step_id")
    op.drop_column("file_uploads", "uploaded_on_branch_id")
    op.drop_column("file_uploads", "content_summary")
    op.drop_column("file_uploads", "extracted_content")
    op.drop_column("file_uploads", "ai_description")

    op.drop_constraint("ck_ai_session_steps_step_type", "ai_session_steps", type_="check")
    op.create_check_constraint(
        "ck_ai_session_steps_step_type", "ai_session_steps",
        "step_type IN ('question', 'action', 'script_generation', 'verification', 'info_request', 'note', 'intake_analysis')",
    )
    op.drop_index("ix_ai_session_steps_branch_id", "ai_session_steps")
    op.drop_column("ai_session_steps", "fork_point_id")
    op.drop_column("ai_session_steps", "is_fork_point")
    op.drop_column("ai_session_steps", "branch_id")

    op.drop_column("ai_sessions", "total_parked_seconds")
    op.drop_column("ai_sessions", "total_active_seconds")
    op.drop_column("ai_sessions", "handoff_count")
    op.drop_column("ai_sessions", "active_branch_id")
    op.drop_column("ai_sessions", "is_branching")

    op.drop_table("session_resolution_outputs")
    op.drop_table("session_handoffs")
    op.drop_table("fork_points")
    op.drop_table("session_branches")