"""add users table and user_id columns

Revision ID: 0008
Revises: 0007
Create Date: 2026-02-11 00:00:00.000000
"""

from alembic import op

revision = "0008"
down_revision = "0007"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Create users table
    op.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id SERIAL PRIMARY KEY,
            username TEXT UNIQUE NOT NULL,
            email TEXT,
            password_hash TEXT NOT NULL,
            role TEXT NOT NULL DEFAULT 'user',
            created_at TIMESTAMPTZ DEFAULT now()
        )
    """)
    op.execute("CREATE INDEX IF NOT EXISTS ix_users_username ON users (username)")

    # Add user_id columns WITHOUT FK constraints first (so we can assign existing rows)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE notes ADD COLUMN user_id INTEGER;
        EXCEPTION WHEN duplicate_column THEN NULL;
        END $$
    """)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE conversations ADD COLUMN user_id INTEGER;
        EXCEPTION WHEN duplicate_column THEN NULL;
        END $$
    """)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE settings ADD COLUMN user_id INTEGER NOT NULL DEFAULT 1;
        EXCEPTION WHEN duplicate_column THEN NULL;
        END $$
    """)

    # Assign all existing data to user_id=1 (the first registered user gets SERIAL id=1)
    op.execute("UPDATE notes SET user_id = 1 WHERE user_id IS NULL")
    op.execute("UPDATE conversations SET user_id = 1 WHERE user_id IS NULL")
    op.execute("UPDATE settings SET user_id = 1 WHERE user_id != 1")

    # Restructure settings PK: drop old single-column PK, create composite (user_id, key)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE settings DROP CONSTRAINT settings_pkey;
        EXCEPTION WHEN undefined_object THEN NULL;
        END $$
    """)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE settings ADD CONSTRAINT settings_pkey PRIMARY KEY (user_id, key);
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$
    """)

    # Now add FK constraints with NOT VALID (skips validation of existing rows
    # where user_id=1 has no matching user yet — becomes valid once first user registers)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE notes ADD CONSTRAINT notes_user_id_fkey
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE NOT VALID;
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$
    """)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE conversations ADD CONSTRAINT conversations_user_id_fkey
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE NOT VALID;
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$
    """)
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE settings ADD CONSTRAINT settings_user_id_fkey
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE NOT VALID;
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$
    """)

    op.execute("CREATE INDEX IF NOT EXISTS ix_notes_user_id ON notes (user_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_conversations_user_id ON conversations (user_id)")


def downgrade() -> None:
    # Reverse settings changes
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE settings DROP CONSTRAINT settings_pkey;
        EXCEPTION WHEN undefined_object THEN NULL;
        END $$
    """)
    op.execute("ALTER TABLE settings ALTER COLUMN user_id DROP NOT NULL")
    op.execute("""
        DO $$ BEGIN
            ALTER TABLE settings ADD CONSTRAINT settings_pkey PRIMARY KEY (key);
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$
    """)
    op.execute("ALTER TABLE settings DROP COLUMN IF EXISTS user_id")

    op.execute("DROP INDEX IF EXISTS ix_conversations_user_id")
    op.execute("ALTER TABLE conversations DROP COLUMN IF EXISTS user_id")

    op.execute("DROP INDEX IF EXISTS ix_notes_user_id")
    op.execute("ALTER TABLE notes DROP COLUMN IF EXISTS user_id")

    op.execute("DROP INDEX IF EXISTS ix_users_username")
    op.execute("DROP TABLE IF EXISTS users")