import os
from dotenv import load_dotenv
import sqlalchemy
from sqlalchemy import create_engine, text

# Load environment variables
load_dotenv('.env')

DATABASE_URL = os.environ.get('DATABASE_URL')
if not DATABASE_URL:
    print("DATABASE_URL not found in .env")
    exit(1)

# Ensure it's a postgresql url
if DATABASE_URL.startswith("postgres://"):
    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)

print(f"Connecting to database: {DATABASE_URL}")
engine = create_engine(DATABASE_URL)

sql_statements = [
    # 1. Create new tables
    """
    CREATE TABLE IF NOT EXISTS public.gerencias_direcciones
    (
        id_gerencias_direcciones integer NOT NULL DEFAULT nextval('gerencias_direcciones_id_gerencias_direcciones_seq'::regclass),
        nombre character varying(500) COLLATE pg_catalog."default",
        valido boolean DEFAULT true,
        fecha_registro time with time zone DEFAULT now(),
        CONSTRAINT gerencias_direcciones_pkey PRIMARY KEY (id_gerencias_direcciones)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS public.gerencias_lineas_divisiones
    (
        id_gerencias_lineas_divisiones integer NOT NULL DEFAULT nextval('gerencias_lineas_divisiones_id_gerencias_lineas_divisiones_seq'::regclass),
        nombres character varying(500) COLLATE pg_catalog."default",
        id_gerencias_direcciones integer,
        valido boolean DEFAULT true,
        fecha_registro time with time zone DEFAULT now(),
        CONSTRAINT gerencias_lineas_divisiones_pkey PRIMARY KEY (id_gerencias_lineas_divisiones)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS public.coordinaciones_divisiones
    (
        id_coordinaciones_divisiones integer NOT NULL DEFAULT nextval('coordinaciones_divisiones_id_coordinaciones_divisiones_seq'::regclass),
        nombres character varying(500) COLLATE pg_catalog."default",
        id_gerencias_lineas_divisiones integer,
        id_gerencias_direcciones integer,
        valido boolean DEFAULT true,
        fecha_registro time with time zone DEFAULT now(),
        CONSTRAINT coordinaciones_divisiones_pkey PRIMARY KEY (id_coordinaciones_divisiones)
    )
    """,
    # Note: Sequences need to be created if they don't exist. The user's SQL assumes they exist. 
    # If they don't, this will fail. Let's create sequences first to be safe, or just use SERIAL.
    # The user provided SQL with specific sequence names. Let's create them if they don't exist.
]

# We should probably fix the sequence issues before running. Let's rewrite the CREATE statements to use SERIAL to automatically handle sequences if this is a fresh setup, or just run the user's exact SQL but wrapped in sequence creation if needed.
# Since the user specifically provided the SQL with `nextval('...')`, it's possible the DB already has these sequences or they were exported from pgAdmin.

# A safer approach for the python script: Let's use SERIAL for the IDs, which creates the sequences automatically in Postgres, BUT since the user provided explicit nextval, maybe they are already created. I will try to create sequences first, ignoring errors if they exist.

setup_sequences_sql = [
    "CREATE SEQUENCE IF NOT EXISTS gerencias_direcciones_id_gerencias_direcciones_seq;",
    "CREATE SEQUENCE IF NOT EXISTS gerencias_lineas_divisiones_id_gerencias_lineas_divisiones_seq;",
    "CREATE SEQUENCE IF NOT EXISTS coordinaciones_divisiones_id_coordinaciones_divisiones_seq;",
    "CREATE SEQUENCE IF NOT EXISTS ubicacion_personal_id_ubicacion_personal_seq;"
]

alter_statements = [
    # 2. Modify ubicacion_personal
    "ALTER TABLE public.ubicacion_personal DROP COLUMN IF EXISTS id_gerencia;",
    "ALTER TABLE public.ubicacion_personal DROP COLUMN IF EXISTS id_gerencia_linea;",
    "ALTER TABLE public.ubicacion_personal DROP COLUMN IF EXISTS id_division;",
    
    "ALTER TABLE public.ubicacion_personal ADD COLUMN IF NOT EXISTS id_gerencias_direcciones integer;",
    "ALTER TABLE public.ubicacion_personal ADD COLUMN IF NOT EXISTS id_gerencias_lineas_divisiones integer;",
    "ALTER TABLE public.ubicacion_personal ADD COLUMN IF NOT EXISTS id_coordinaciones_divisiones integer;",
    
    # Add foreign keys (optional but good practice)
    # "ALTER TABLE public.ubicacion_personal ADD CONSTRAINT fk_gerencias_direcciones FOREIGN KEY (id_gerencias_direcciones) REFERENCES public.gerencias_direcciones(id_gerencias_direcciones);",
    
    # 3. Modify odi
    "ALTER TABLE public.odi DROP COLUMN IF EXISTS id_gerencia;",
    "ALTER TABLE public.odi DROP COLUMN IF EXISTS id_gerencia_linea;",
    "ALTER TABLE public.odi DROP COLUMN IF EXISTS id_division;"
]

# Note: We are not dropping the old tables (direccion, gerencia, gerencia_linea, division) yet just in case they have other dependencies we missed, but they are no longer used by our models.

try:
    with engine.connect() as conn:
        for seq in setup_sequences_sql:
            try:
                conn.execute(text(seq))
            except Exception as e:
                print(f"Sequence info: {e}") # might already exist
        
        print("Creating tables...")
        for sql in sql_statements:
            conn.execute(text(sql))
            
        print("Altering existing tables...")
        for sql in alter_statements:
            try:
                conn.execute(text(sql))
            except Exception as e:
                print(f"Alter info: {e}")
                
        conn.commit()
        print("Database update complete!")
except Exception as e:
    print(f"Error updating database: {e}")