"""Add workspaces table, workspace_id to trees and categories, color to categories

Revision ID: 036
Revises: 035
Create Date: 2026-02-15

Adds workspace system for organizational context above folders.
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID


# revision identifiers, used by Alembic.
revision: str = '036'
down_revision: Union[str, None] = '035'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Create workspaces table
    op.create_table(
        'workspaces',
        sa.Column('id', UUID(as_uuid=True), primary_key=True, server_default=sa.text('gen_random_uuid()')),
        sa.Column('name', sa.String(100), nullable=False),
        sa.Column('slug', sa.String(100), nullable=False),
        sa.Column('description', sa.Text(), nullable=True),
        sa.Column('icon', sa.String(10), nullable=True),
        sa.Column('accent_color', sa.String(7), nullable=True),
        sa.Column('account_id', UUID(as_uuid=True), sa.ForeignKey('accounts.id', ondelete='CASCADE'), nullable=False),
        sa.Column('is_default', sa.Boolean(), nullable=False, server_default='false'),
        sa.Column('sort_order', sa.Integer(), nullable=False, server_default='0'),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()')),
        sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()')),
        sa.UniqueConstraint('slug', 'account_id', name='uq_workspaces_slug_account'),
    )
    op.create_index('ix_workspaces_slug', 'workspaces', ['slug'])
    op.create_index('ix_workspaces_account_id', 'workspaces', ['account_id'])

    # Add workspace_id to trees
    op.add_column('trees', sa.Column('workspace_id', UUID(as_uuid=True), nullable=True))
    op.create_foreign_key('fk_trees_workspace_id', 'trees', 'workspaces', ['workspace_id'], ['id'], ondelete='SET NULL')
    op.create_index('ix_trees_workspace_id', 'trees', ['workspace_id'])

    # Add color and workspace_id to tree_categories
    op.add_column('tree_categories', sa.Column('color', sa.String(7), nullable=True, server_default='#3b82f6'))
    op.add_column('tree_categories', sa.Column('workspace_id', UUID(as_uuid=True), nullable=True))
    op.create_foreign_key('fk_tree_categories_workspace_id', 'tree_categories', 'workspaces', ['workspace_id'], ['id'], ondelete='SET NULL')
    op.create_index('ix_tree_categories_workspace_id', 'tree_categories', ['workspace_id'])

    # Seed default workspaces for each existing account and assign trees
    op.execute("""
        -- Create default workspaces for each account
        INSERT INTO workspaces (name, slug, description, icon, accent_color, account_id, is_default, sort_order)
        SELECT 'Troubleshooting', 'troubleshooting', 'Break/fix decision trees', '🔧', '#ef4444', a.id, true, 0
        FROM accounts a;

        INSERT INTO workspaces (name, slug, description, icon, accent_color, account_id, is_default, sort_order)
        SELECT 'Procedures', 'procedures', 'Step-by-step operational flows', '📋', '#3b82f6', a.id, false, 1
        FROM accounts a;

        INSERT INTO workspaces (name, slug, description, icon, accent_color, account_id, is_default, sort_order)
        SELECT 'Policies', 'policies', 'Compliance & policy builders', '📜', '#8b5cf6', a.id, false, 2
        FROM accounts a;

        INSERT INTO workspaces (name, slug, description, icon, accent_color, account_id, is_default, sort_order)
        SELECT 'Finance', 'finance', 'Billing & procurement flows', '💰', '#22c55e', a.id, false, 3
        FROM accounts a;

        -- Assign existing trees to appropriate workspace based on tree_type
        UPDATE trees t
        SET workspace_id = w.id
        FROM workspaces w
        WHERE w.account_id = t.account_id
          AND w.slug = 'troubleshooting'
          AND t.tree_type = 'troubleshooting';

        UPDATE trees t
        SET workspace_id = w.id
        FROM workspaces w
        WHERE w.account_id = t.account_id
          AND w.slug = 'procedures'
          AND t.tree_type = 'procedural';
    """)


def downgrade() -> None:
    op.drop_index('ix_tree_categories_workspace_id', 'tree_categories')
    op.drop_constraint('fk_tree_categories_workspace_id', 'tree_categories', type_='foreignkey')
    op.drop_column('tree_categories', 'workspace_id')
    op.drop_column('tree_categories', 'color')

    op.drop_index('ix_trees_workspace_id', 'trees')
    op.drop_constraint('fk_trees_workspace_id', 'trees', type_='foreignkey')
    op.drop_column('trees', 'workspace_id')

    op.drop_index('ix_workspaces_account_id', 'workspaces')
    op.drop_index('ix_workspaces_slug', 'workspaces')
    op.drop_table('workspaces')