#!/usr/bin/env python3
import argparse
import asyncio
import csv
import os
import sys
import random
import threading
import re
from typing import List, Dict, Any
import psycopg2
from psycopg2.errors import (
    UniqueViolation,
    UndefinedTable,
    UndefinedColumn,
    DuplicateObject,
)
from supabase import create_client, Client
from colorama import init, Fore

init(autoreset=True)

SOURCE_SUPABASE_URL = "https://x.supabase.co"
SOURCE_SUPABASE_ANON_KEY = "x"

TARGET_SUPABASE_URL = "https://athena-db.com"
TARGET_SUPABASE_ANON_KEY = "x"

csv_lock = threading.Lock()


def create_supabase_client(url, key) -> Client:
    return create_client(url, key)


def connect_postgres(pg_url: str):
    return psycopg2.connect(pg_url)


def create_anon_role_if_needed(pg_url: str):
    """
    Create the 'anon' role in Postgres if it does not exist.
    This is to prevent 'role "anon" does not exist' errors when running views.sql or others.
    """
    try:
        conn = connect_postgres(pg_url)
        cur = conn.cursor()
        # Check if 'anon' role exists
        cur.execute("SELECT 1 FROM pg_roles WHERE rolname = 'anon'")
        exists = cur.fetchone()
        if not exists:
            print(f"{Fore.YELLOW}Role 'anon' does not exist, creating...")
            try:
                cur.execute("CREATE ROLE anon")
                conn.commit()
                print(f"{Fore.GREEN}Role 'anon' created.")
            except DuplicateObject:
                print(f"{Fore.BLUE}Role 'anon' already exists (race condition).")
            except Exception as ce:
                print(f"{Fore.RED}Failed to create 'anon' role: {ce}")
                conn.rollback()
        else:
            print(f"{Fore.CYAN}Role 'anon' already exists.")
        cur.close()
        conn.close()
    except Exception as e:
        print(f"{Fore.RED}Error checking/creating 'anon' role: {e}")


def execute_sql_file(pg_url: str, path: str):
    if not os.path.exists(path):
        print(f"{Fore.YELLOW}Skipping missing file: {path}")
        return
    with open(path, "r", encoding="utf-8") as f:
        sql = f.read()
    conn = connect_postgres(pg_url)
    cur = conn.cursor()
    try:
        cur.execute(sql)
        conn.commit()
        print(f"{Fore.GREEN}Executed {path}")
    except Exception as e:
        # Intercept error for missing role "anon" and create it, then retry ONCE
        if 'role "anon" does not exist' in str(e):
            print(f"{Fore.YELLOW}Detected error: {e}")
            print(
                f"{Fore.CYAN}Attempting to create role 'anon' and retry executing {path} ..."
            )
            try:
                create_anon_role_if_needed(pg_url)
                # Retry once
                try:
                    cur.execute(sql)
                    conn.commit()
                    print(f"{Fore.GREEN}Executed {path} after creating 'anon' role")
                    return
                except Exception as exp2:
                    print(f"{Fore.RED}Still failed executing {path}: {exp2}")
                    conn.rollback()
            except Exception as ce2:
                print(
                    f"{Fore.RED}Failed to create 'anon' role or re-execute file: {ce2}"
                )
                conn.rollback()
        else:
            print(f"{Fore.RED}Error executing {path}: {e}")
            conn.rollback()
    finally:
        cur.close()
        conn.close()


def add_missing_column(conn, table, column):
    """
    Attempt to add the missing column to the target table.
    For this script, we add as TEXT and nullable. In real migrations, inspect or map types as needed,
    or use source row sample to infer.
    """
    cur = conn.cursor()
    try:
        print(
            f"{Fore.YELLOW}Attempting to add missing column '{column}' to table '{table}' as TEXT..."
        )
        alter_sql = f'ALTER TABLE "{table}" ADD COLUMN IF NOT EXISTS "{column}" TEXT;'
        cur.execute(alter_sql)
        conn.commit()
        print(f"{Fore.GREEN}Added column '{column}' to '{table}'.")
    except Exception as e:
        print(f"{Fore.RED}Could not add column '{column}' to table '{table}': {e}")
        conn.rollback()
    finally:
        cur.close()


def print_table_schema(pg_url, table="users", schema="public"):
    try:
        conn = connect_postgres(pg_url)
        cur = conn.cursor()
        cur.execute(
            """
            SELECT column_name, data_type, is_nullable, column_default
            FROM information_schema.columns
            WHERE table_schema = %s AND table_name = %s
            ORDER BY ordinal_position
        """,
            (schema, table),
        )
        columns = cur.fetchall()
        if not columns:
            print(f"{Fore.YELLOW}No columns found for public.{table}")
        else:
            print(f"{Fore.CYAN}Schema for public.{table}:")
            for col in columns:
                colname, datatype, isnull, coldefault = col
                print(
                    f"  {colname} {datatype} {'NULL' if isnull == 'YES' else 'NOT NULL'} DEFAULT {coldefault}"
                )
        cur.close()
        conn.close()
    except Exception as ex:
        print(f"{Fore.RED}Failed to print schema for public.{table}: {ex}")


async def retry_with_backoff(func, max_retries=3, base_delay=5.0, *args, **kwargs):
    allow_broken_columns = kwargs.pop("allow_broken_columns", False)
    original_rows = kwargs.get("rows", None)
    already_removed = kwargs.pop("already_removed", set())
    last_exception = None
    pg_url = kwargs.get("_pg_url")
    table = kwargs.get("table")

    for attempt in range(max_retries):
        try:
            return await func(*args, **kwargs)
        except Exception as e:
            last_exception = e
            err_code = None
            broken_column = None
            err_msg = str(e)
            err_msg_full = err_msg

            # Try to extract error info from supabase response dict in e.args[0]
            if hasattr(e, "args") and e.args and isinstance(e.args[0], dict):
                err_code = e.args[0].get("code")
                err_msg_dict = e.args[0]
                err_msg_full = str(err_msg_dict)
                msg = err_msg_dict.get("message", "")
                m = re.search(r"Could not find the '(\w+)' column", msg)
                if m:
                    broken_column = m.group(1)
            # Print schema if "Could not find" for table "users"
            if (
                ("Could not find the" in err_msg or err_code == "PGRST204")
                and table == "users"
                and pg_url
            ):
                print(
                    f"{Fore.YELLOW}Attempting to print schema for public.users due to missing column or similar error."
                )
                print_table_schema(pg_url, table="users", schema="public")
            # For Supabase: try to add column to Postgres too, if using direct PG connection
            if "Could not find the" in err_msg or err_code == "PGRST204":
                if pg_url and table and broken_column:
                    try:
                        conn = connect_postgres(pg_url)
                        add_missing_column(conn, table, broken_column)
                        conn.close()
                        # Try again immediately (do not count toward retries)
                        continue
                    except Exception as addcol_exc:
                        print(
                            f"{Fore.RED}Tried to fix table by adding column but failed: {addcol_exc}"
                        )
            # Broken columns logic (if allow_broken_columns, fallback to dropping)
            if (
                "Could not find the" in err_msg or err_code == "PGRST204"
            ) and allow_broken_columns:
                if original_rows is not None and broken_column:
                    if broken_column in already_removed:
                        print(
                            f"{Fore.RED}Tried to remove already removed column '{broken_column}'. Giving up."
                        )
                        break
                    print(
                        f"{Fore.YELLOW}Removing missing column '{broken_column}' and retrying insert..."
                    )
                    rows_no_col = []
                    for row in original_rows:
                        row2 = dict(row)
                        row2.pop(broken_column, None)
                        rows_no_col.append(row2)
                    if not rows_no_col or not rows_no_col[0]:
                        print(
                            f"{Fore.RED}Cannot proceed: All columns are missing after removing '{broken_column}'."
                        )
                        break
                    new_kwargs = dict(kwargs)
                    new_kwargs["rows"] = rows_no_col
                    new_kwargs["allow_broken_columns"] = True
                    new_kwargs["already_removed"] = set(already_removed)
                    new_kwargs["already_removed"].add(broken_column)
                    if pg_url:
                        new_kwargs["_pg_url"] = pg_url
                    return await retry_with_backoff(
                        func,
                        max_retries=max_retries,
                        base_delay=base_delay,
                        *args,
                        **new_kwargs,
                    )
                else:
                    print(f"{Fore.RED}Schema error (missing column): {err_msg_full}")
                    break
            elif "Could not find the" in err_msg or err_code == "PGRST204":
                print(f"{Fore.RED}Schema error: {err_msg_full}")
                break
            delay = base_delay * (2**attempt) + random.uniform(0, 2)
            print(
                f"{Fore.YELLOW}Retry {attempt+1}/{max_retries}: {err_msg} ({delay:.1f}s)"
            )
            await asyncio.sleep(delay)
    if last_exception:
        print(
            f"{Fore.RED}Operation failed after {max_retries} attempts: {last_exception}"
        )
        raise last_exception
    raise Exception(f"Max retries ({max_retries}) exceeded")


def save_to_csv(table, rows):
    if not rows:
        return
    os.makedirs("./scripts/data", exist_ok=True)
    csv_path = f"./scripts/data/{table}.csv"
    with csv_lock:
        file_exists = os.path.isfile(csv_path)
        with open(csv_path, "a", newline="", encoding="utf-8") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=rows[0].keys())
            if not file_exists:
                writer.writeheader()
            writer.writerows(rows)


async def fetch_batch_supabase(client, table, last_id, batch_size):
    try:
        query = (
            client.table(table)
            .select("*")
            .gt("id", last_id)
            .order("id")
            .limit(batch_size)
            .execute()
        )
        return query.data
    except Exception as e:
        print(f"{Fore.RED}Error fetching batch from table '{table}': {e}")
        raise


async def insert_batch_supabase(client, table, rows):
    try:
        client.table(table).upsert(rows).execute()
    except Exception as e:
        if table == "users":
            pg_url = os.getenv("TARGET_POSTGRES_URL")
            if pg_url:
                print(
                    f"{Fore.YELLOW}Printing schema for public.users due to insert error."
                )
                print_table_schema(pg_url, table="users", schema="public")
            else:
                print(
                    f"{Fore.YELLOW}TARGET_POSTGRES_URL not set; can't print schema for users."
                )
        print(f"{Fore.RED}Error inserting batch into table '{table}': {e}")
        raise


async def fetch_and_insert_table(
    source_client, target_client, table, batch_size, allow_broken_columns=False
):
    last_id = 0
    while True:
        try:
            rows = await retry_with_backoff(
                fetch_batch_supabase,
                table=table,
                last_id=last_id,
                batch_size=batch_size,
                client=source_client,
            )
        except Exception as e:
            print(f"{Fore.RED}Migration stopped during fetch: {e}")
            break
        if not rows:
            print(f"{Fore.CYAN}No more rows in {table}")
            break
        save_to_csv(table, rows)
        try:
            await retry_with_backoff(
                insert_batch_supabase,
                table=table,
                rows=rows,
                client=target_client,
                allow_broken_columns=allow_broken_columns,
            )
        except Exception as e:
            if table == "users":
                pg_url = os.getenv("TARGET_POSTGRES_URL")
                if pg_url:
                    print(
                        f"{Fore.YELLOW}Printing schema for public.users due to insert error."
                    )
                    print_table_schema(pg_url, table="users", schema="public")
                else:
                    print(
                        f"{Fore.YELLOW}TARGET_POSTGRES_URL not set; can't print schema for users."
                    )
            print(f"{Fore.RED}Migration stopped during insert: {e}")
            break
        last_id = rows[-1].get("id", last_id)
        print(f"{Fore.GREEN}Migrated {len(rows)} rows from {table}, last_id={last_id}")


async def migrate_table(args):
    if args.pg_url:
        print(f"{Fore.CYAN}Running schema SQL files on PostgreSQL...")
        # Before running SQL, ensure anon role exists to prevent 'role "anon" does not exist' errors
        create_anon_role_if_needed(args.pg_url)
        for file in ("seed.sql", "views.sql", "extensions.sql"):
            try:
                execute_sql_file(args.pg_url, file)
            except Exception as e:
                print(f"{Fore.RED}Could not execute {file} on Postgres: {e}")
                return

    source_client = create_supabase_client(
        SOURCE_SUPABASE_URL, SOURCE_SUPABASE_ANON_KEY
    )
    if args.pg_url:
        target_client = None
    else:
        target_client = create_supabase_client(
            TARGET_SUPABASE_URL, TARGET_SUPABASE_ANON_KEY
        )

    allow_broken_columns = getattr(args, "allow_broken_columns", False)

    if target_client:
        await fetch_and_insert_table(
            source_client,
            target_client,
            args.table,
            args.batch_size,
            allow_broken_columns=allow_broken_columns,
        )
    else:
        conn = connect_postgres(args.pg_url)
        cur = conn.cursor()
        last_id = 0
        broken_columns = set()
        while True:
            try:
                rows = await retry_with_backoff(
                    fetch_batch_supabase,
                    table=args.table,
                    last_id=last_id,
                    batch_size=args.batch_size,
                    client=source_client,
                    _pg_url=args.pg_url,
                )
            except Exception as e:
                print(f"{Fore.RED}Migration stopped during fetch: {e}")
                break
            if not rows:
                print(f"{Fore.CYAN}No more rows in {args.table}")
                break
            save_to_csv(args.table, rows)
            insert_rows = rows[:]
            while True:
                try:
                    if broken_columns:
                        insert_rows = [
                            {k: v for k, v in row.items() if k not in broken_columns}
                            for row in insert_rows
                        ]
                    if not insert_rows or not insert_rows[0]:
                        print(
                            f"{Fore.RED}All columns are missing or removed, cannot proceed."
                        )
                        break
                    cols = list(insert_rows[0].keys())
                    vals = [[row[c] for c in cols] for row in insert_rows]
                    placeholders = ", ".join(["%s"] * len(cols))
                    colnames = ", ".join(cols)
                    query = f'INSERT INTO "{args.table}" ({colnames}) VALUES ({placeholders}) ON CONFLICT DO NOTHING'
                    cur.executemany(query, vals)
                    conn.commit()
                    if broken_columns:
                        print(
                            f"{Fore.GREEN}Inserted {len(insert_rows)} rows into {args.table} (after removing columns: {', '.join(broken_columns)})"
                        )
                    else:
                        print(
                            f"{Fore.GREEN}Inserted {len(insert_rows)} rows into {args.table}"
                        )
                    break
                except UndefinedTable as ute:
                    print(f"{Fore.RED}Table not found: {ute}")
                    conn.rollback()
                    break
                except UndefinedColumn as uc:
                    uc_msg = str(uc)
                    m = re.search(r'column "([^"]+)" does not exist', uc_msg)
                    if m:
                        missing_col = m.group(1)
                        print(
                            f"{Fore.YELLOW}Column '{missing_col}' does not exist in '{args.table}', attempting to add..."
                        )
                        add_missing_column(conn, args.table, missing_col)
                        continue
                    if allow_broken_columns:
                        print(
                            f"{Fore.YELLOW}UndefinedColumn fallback, removing unknown column: {uc}"
                        )
                        unknown_col = getattr(uc, "diag", None) or "unknown"
                        broken_columns.add(str(unknown_col))
                        conn.rollback()
                        continue
                    else:
                        print(f"{Fore.RED}Column error: {uc}")
                        conn.rollback()
                        break
                except UniqueViolation:
                    print(f"{Fore.YELLOW}Duplicate key encountered, continuing.")
                    conn.rollback()
                    break
                except Exception as e:
                    if args.table == "users" and args.pg_url:
                        print(
                            f"{Fore.YELLOW}Printing schema for public.users due to error inserting rows."
                        )
                        print_table_schema(args.pg_url, table="users", schema="public")
                    print(f"{Fore.RED}Error inserting rows: {e}")
                    conn.rollback()
                    break
            if not insert_rows or not insert_rows[0]:
                break
            last_id = rows[-1].get("id", last_id)
        cur.close()
        conn.close()


def main():
    parser = argparse.ArgumentParser(
        description="Migrate Supabase data or run schema SQLs"
    )
    parser.add_argument("--table", required=True, help="Table name to migrate")
    parser.add_argument("--batch-size", type=int, default=500, help="Rows per batch")
    parser.add_argument(
        "--pg-url", help="Optional PostgreSQL URL (postgresql://user:pass@host/db)"
    )
    parser.add_argument(
        "--allow-broken-columns",
        action="store_true",
        help="If set, drops columns missing from schema when inserting batches.",
    )
    args = parser.parse_args()
    try:
        asyncio.run(migrate_table(args))
    except KeyboardInterrupt:
        print(f"{Fore.YELLOW}\nAborted by user.")
    except Exception as e:
        print(f"{Fore.RED}Fatal error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
