"""add ai_sessions and ai_session_steps tables

Revision ID: f1a2b3c4d5e6
Revises: ee98013dd18c
Create Date: 2026-03-18
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID, JSONB

# revision identifiers, used by Alembic.
revision = "f1a2b3c4d5e6"
down_revision = "ee98013dd18c"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ── ai_sessions table ──
    op.create_table(
        "ai_sessions",
        sa.Column("id", UUID(as_uuid=True), primary_key=True),
        sa.Column("user_id", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True),
        sa.Column("account_id", UUID(as_uuid=True), sa.ForeignKey("accounts.id", ondelete="CASCADE"), nullable=False, index=True),
        sa.Column("team_id", UUID(as_uuid=True), sa.ForeignKey("teams.id", ondelete="SET NULL"), nullable=True, index=True),
        # Intake
        sa.Column("intake_type", sa.String(20), nullable=False, server_default="free_text"),
        sa.Column("intake_content", JSONB, nullable=False, server_default="{}"),
        sa.Column("problem_summary", sa.Text, nullable=True),
        sa.Column("problem_domain", sa.String(100), nullable=True),
        # Session state
        sa.Column("status", sa.String(20), nullable=False, server_default="active", index=True),
        sa.Column("confidence_tier", sa.String(20), nullable=False, server_default="discovery"),
        sa.Column("confidence_score", sa.Float, nullable=False, server_default="0.0"),
        # Flow matching
        sa.Column("matched_flow_id", UUID(as_uuid=True), sa.ForeignKey("trees.id", ondelete="SET NULL"), nullable=True),
        sa.Column("match_score", sa.Float, nullable=True),
        # PSA link
        sa.Column("psa_ticket_id", sa.String(100), nullable=True),
        sa.Column("psa_connection_id", UUID(as_uuid=True), sa.ForeignKey("psa_connections.id", ondelete="SET NULL"), nullable=True),
        sa.Column("ticket_data", JSONB, nullable=True),
        # Resolution / Escalation
        sa.Column("resolution_summary", sa.Text, nullable=True),
        sa.Column("resolution_action", sa.Text, nullable=True),
        sa.Column("escalation_reason", sa.Text, nullable=True),
        sa.Column("escalation_package", JSONB, nullable=True),
        sa.Column("escalated_to_id", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=True),
        # Feedback
        sa.Column("session_rating", sa.Integer, nullable=True),
        sa.Column("session_feedback", sa.Text, nullable=True),
        # AI tracking
        sa.Column("total_input_tokens", sa.Integer, nullable=False, server_default="0"),
        sa.Column("total_output_tokens", sa.Integer, nullable=False, server_default="0"),
        sa.Column("step_count", sa.Integer, nullable=False, server_default="0"),
        # Timestamps
        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.Column("resolved_at", sa.DateTime(timezone=True), nullable=True),
        # LLM context
        sa.Column("system_prompt_snapshot", sa.Text, nullable=True),
        sa.Column("conversation_messages", JSONB, nullable=False, server_default="[]"),
        # Check constraints
        sa.CheckConstraint(
            "intake_type IN ('free_text', 'psa_ticket', 'screenshot', 'log_paste', 'combined')",
            name="ck_ai_sessions_intake_type",
        ),
        sa.CheckConstraint(
            "status IN ('active', 'paused', 'resolved', 'escalated', 'abandoned')",
            name="ck_ai_sessions_status",
        ),
        sa.CheckConstraint(
            "confidence_tier IN ('guided', 'exploring', 'discovery')",
            name="ck_ai_sessions_confidence_tier",
        ),
    )

    # ── ai_session_steps table ──
    op.create_table(
        "ai_session_steps",
        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, index=True),
        sa.Column("step_order", sa.Integer, nullable=False),
        sa.Column("step_type", sa.String(30), nullable=False),
        # Content
        sa.Column("content", JSONB, nullable=False, server_default="{}"),
        sa.Column("context_message", sa.Text, nullable=True),
        # Options
        sa.Column("options_presented", JSONB, nullable=True),
        # Engineer response
        sa.Column("selected_option", sa.String(500), nullable=True),
        sa.Column("free_text_input", sa.Text, nullable=True),
        sa.Column("was_free_text", sa.Boolean, nullable=False, server_default="false"),
        sa.Column("was_skipped", sa.Boolean, nullable=False, server_default="false"),
        # Action results
        sa.Column("action_result", JSONB, nullable=True),
        # Script generation link
        sa.Column("script_generation_id", UUID(as_uuid=True), sa.ForeignKey("script_generations.id", ondelete="SET NULL"), nullable=True),
        # AI internals
        sa.Column("confidence_at_step", sa.Float, nullable=False, server_default="0.0"),
        sa.Column("ai_reasoning", sa.Text, nullable=True),
        sa.Column("input_tokens", sa.Integer, nullable=False, server_default="0"),
        sa.Column("output_tokens", sa.Integer, nullable=False, server_default="0"),
        # Timestamps
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("responded_at", sa.DateTime(timezone=True), nullable=True),
        # Check constraint
        sa.CheckConstraint(
            "step_type IN ('question', 'action', 'script_generation', 'verification', "
            "'info_request', 'note', 'intake_analysis')",
            name="ck_ai_session_steps_step_type",
        ),
    )

    # ── Add flow matching columns to trees table ──
    op.add_column("trees", sa.Column("origin", sa.String(20), nullable=True, comment="manual | ai_generated | ai_enhanced"))
    op.add_column("trees", sa.Column("source_session_id", UUID(as_uuid=True), nullable=True))
    op.add_column("trees", sa.Column("match_keywords", JSONB, nullable=True, server_default="[]"))
    op.add_column("trees", sa.Column("success_rate", sa.Float, nullable=True))
    op.add_column("trees", sa.Column("last_matched_at", sa.DateTime(timezone=True), nullable=True))


def downgrade() -> None:
    op.drop_column("trees", "last_matched_at")
    op.drop_column("trees", "success_rate")
    op.drop_column("trees", "match_keywords")
    op.drop_column("trees", "source_session_id")
    op.drop_column("trees", "origin")
    op.drop_table("ai_session_steps")
    op.drop_table("ai_sessions")
