import os
import asyncio
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
import pandas as pd
from dotenv import load_dotenv

# Try to find .env in current dir and Agent dir
load_dotenv()
load_dotenv(dotenv_path=os.path.join(os.path.dirname(__file__), '.env'))
load_dotenv(dotenv_path=os.path.join(os.path.dirname(__file__), '..', '.env'))

# LOCAL_DB = "postgresql://postgres:postgres@localhost:5432/employee"
LOCAL_DB = os.environ.get("DB_CONNECTION_URL")
# NEON_DB = Your new Neon connection string
NEON_DB = os.environ.get("NEON_DB_URL")

async def migrate():
    if not NEON_DB:
        print("❌ Error: NEON_DB_URL not found in environment variables.")
        return

    print("🚀 Starting Cloud Migration to Neon.tech...")
    
    # 1. Connect to Local
    local_engine = create_engine(LOCAL_DB)
    
    # 2. Connect to Neon (Ensure SSL)
    neon_uri = NEON_DB
    if 'sslmode=' not in neon_uri:
        sep = '&' if '?' in neon_uri else '?'
        neon_uri += f"{sep}sslmode=require"
    
    neon_engine = create_engine(neon_uri)

    tables = ["department", "employee", "dept_emp", "dept_manager", "salary", "title"]

    for table in tables:
        print(f"📦 Migrating table: {table}...")
        try:
            # Read data from local
            df = pd.read_sql_table(table, local_engine)
            print(f"   - Found {len(df)} rows.")
            
            # Write to Neon (replace if exists)
            df.to_sql(table, neon_engine, if_exists='replace', index=False)
            print(f"   ✅ {table} successfully uploaded to Neon.")
        except Exception as e:
            print(f"   ❌ Failed to migrate {table}: {e}")

    print("\n🎉 Migration Complete! Your cloud database is ready.")

if __name__ == "__main__":
    asyncio.run(migrate())