#!/usr/bin/env python3
"""Migration: add owner_id columns and backfill existing rows.

This script adds an ``owner_id`` column (FK to ``users``) to the
``llm_messages``, ``jobs``, and ``tracked_companies`` tables, then
backfills all existing rows with ``owner_id = 1`` (the default admin user).

It also replaces the old global UNIQUE constraint on
``tracked_companies.company_name`` with a per-owner unique index so that
different users can independently track the same company.

Usage:
    python scripts/migrate_add_owner_id.py

The script is idempotent — running it multiple times is safe.
"""

import os
import sys

import psycopg2

DATABASE_URL = os.getenv(
    "DATABASE_URL",
    "postgresql://postgres:postgres@localhost:5432/sparc",
)

DEFAULT_OWNER_ID = 1


def run_migration():
    """Execute the migration."""
    conn = psycopg2.connect(DATABASE_URL)
    conn.autocommit = False

    try:
        with conn.cursor() as cur:
            # ---------- 1. Add owner_id columns if missing ----------
            cur.execute("""
                DO $$
                BEGIN
                    IF NOT EXISTS (
                        SELECT 1 FROM information_schema.columns
                        WHERE table_name = 'llm_messages' AND column_name = 'owner_id'
                    ) THEN
                        ALTER TABLE llm_messages ADD COLUMN owner_id INTEGER REFERENCES users(id);
                    END IF;

                    IF NOT EXISTS (
                        SELECT 1 FROM information_schema.columns
                        WHERE table_name = 'jobs' AND column_name = 'owner_id'
                    ) THEN
                        ALTER TABLE jobs ADD COLUMN owner_id INTEGER REFERENCES users(id);
                    END IF;

                    IF NOT EXISTS (
                        SELECT 1 FROM information_schema.columns
                        WHERE table_name = 'tracked_companies' AND column_name = 'owner_id'
                    ) THEN
                        ALTER TABLE tracked_companies ADD COLUMN owner_id INTEGER REFERENCES users(id);
                    END IF;
                END $$;
            """)

            # ---------- 2. Backfill owner_id = DEFAULT_OWNER_ID ----------
            cur.execute(
                "UPDATE llm_messages SET owner_id = %s WHERE owner_id IS NULL",
                (DEFAULT_OWNER_ID,),
            )
            messages_updated = cur.rowcount
            print(f"  llm_messages: backfilled {messages_updated} rows")

            cur.execute(
                "UPDATE jobs SET owner_id = %s WHERE owner_id IS NULL",
                (DEFAULT_OWNER_ID,),
            )
            jobs_updated = cur.rowcount
            print(f"  jobs: backfilled {jobs_updated} rows")

            cur.execute(
                "UPDATE tracked_companies SET owner_id = %s WHERE owner_id IS NULL",
                (DEFAULT_OWNER_ID,),
            )
            tracked_updated = cur.rowcount
            print(f"  tracked_companies: backfilled {tracked_updated} rows")

            # ---------- 3. Create indexes ----------
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_messages_owner
                ON llm_messages(owner_id)
            """)
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_jobs_owner
                ON jobs(owner_id)
            """)
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_tracked_companies_owner
                ON tracked_companies(owner_id)
            """)

            # ---------- 4. Replace unique constraint on tracked_companies ----------
            cur.execute("""
                DO $$
                BEGIN
                    IF EXISTS (
                        SELECT 1 FROM pg_constraint
                        WHERE conname = 'tracked_companies_company_name_key'
                    ) THEN
                        ALTER TABLE tracked_companies
                            DROP CONSTRAINT tracked_companies_company_name_key;
                    END IF;
                END $$;
            """)
            cur.execute("""
                CREATE UNIQUE INDEX IF NOT EXISTS uq_tracked_company_owner
                ON tracked_companies(LOWER(company_name), owner_id)
            """)

        conn.commit()
        print("Migration completed successfully.")

    except Exception:
        conn.rollback()
        print("Migration FAILED — rolled back.", file=sys.stderr)
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    print(f"Running owner_id migration against {DATABASE_URL.split('@')[-1]} ...")
    run_migration()