import os
import argparse
import pandas as pd
from sqlalchemy import create_engine, text

def load_database():
    parser = argparse.ArgumentParser(description="Load cleaned supply chain data into MySQL.")
    parser.add_argument("--host", default="localhost", help="MySQL host (default: localhost)")
    parser.add_argument("--port", default="3306", help="MySQL port (default: 3306)")
    parser.add_argument("--user", default="root", help="MySQL username (default: root)")
    parser.add_argument("--password", default="", help="MySQL password (default: empty)")
    args = parser.parse_args()

    cleaned_csv_path = os.path.join("Data", "cleaned_data.csv")
    schema_sql_path = os.path.join("SQL", "schema.sql")

    if not os.path.exists(cleaned_csv_path):
        print(f"Error: Cleaned data file not found at {cleaned_csv_path}. Please run data_cleaning.py first.")
        return

    if not os.path.exists(schema_sql_path):
        print(f"Error: Schema SQL file not found at {schema_sql_path}.")
        return

    # 1. Establish connection to MySQL (without selecting DB first to create database)
    conn_url_base = f"mysql+mysqlconnector://{args.user}:{args.password}@{args.host}:{args.port}"
    print(f"Connecting to MySQL server at {args.host}:{args.port}...")
    
    try:
        engine = create_engine(conn_url_base)
        # Create database first
        with engine.connect() as conn:
            conn.execute(text("CREATE DATABASE IF NOT EXISTS supply_chain_db;"))
            conn.commit()
        print("Database 'supply_chain_db' checked/created.")
        
        # 2. Re-connect with database selected to run the rest of the schema
        conn_url_db = f"{conn_url_base}/supply_chain_db"
        db_engine = create_engine(conn_url_db)
        
        print(f"Executing schema from {schema_sql_path}...")
        with open(schema_sql_path, "r", encoding="utf-8") as f:
            schema_sql = f.read()
        
        # Split sql file into individual commands by semicolon
        sql_commands = [cmd.strip() for cmd in schema_sql.split(";") if cmd.strip()]
        
        with db_engine.connect() as conn:
            # We run the commands one by one
            for cmd in sql_commands:
                # Skip comments, create database and USE statements (since they are already handled)
                if cmd.startswith("--") or cmd.startswith("/*"):
                    continue
                clean_cmd = cmd.upper()
                if clean_cmd.startswith("CREATE DATABASE") or clean_cmd.startswith("USE "):
                    continue
                conn.execute(text(cmd))
                conn.commit()
                
        print("Database schema created successfully.")
    except Exception as e:
        print(f"\nFailed to connect or set up database: {e}")
        print("Please check your MySQL credentials. You can pass them as arguments:")
        print("python Scripts/database_loader.py --user <username> --password <password> --host <host> --port <port>")
        return

    # 3. Connect to the specific supply_chain_db database
    conn_url_db = f"{conn_url_base}/supply_chain_db"
    db_engine = create_engine(conn_url_db)

    # 4. Load cleaned data into memory
    print(f"Reading cleaned data from {cleaned_csv_path}...")
    df = pd.read_csv(cleaned_csv_path)

    # 5. Extract and load normalized tables
    print("Normalizing data and uploading to tables...")

    try:
        # A. Departments
        print("- Loading 'departments'...")
        departments_df = df[["department_id", "department_name"]].drop_duplicates().dropna(subset=["department_id"])
        departments_df.to_sql("departments", con=db_engine, if_exists="append", index=False)
        print(f"  Loaded {len(departments_df)} departments.")

        # B. Categories
        print("- Loading 'categories'...")
        categories_df = df[["category_id", "category_name"]].drop_duplicates().dropna(subset=["category_id"])
        categories_df.to_sql("categories", con=db_engine, if_exists="append", index=False)
        print(f"  Loaded {len(categories_df)} categories.")

        # C. Customers
        print("- Loading 'customers'...")
        # Since customers can have multiple order records, drop duplicates by customer_id
        customers_df = df[[
            "customer_id", "customer_fname", "customer_lname", 
            "customer_segment", "customer_street", "customer_city", 
            "customer_state", "customer_country", "customer_zipcode"
        ]].drop_duplicates(subset=["customer_id"])
        customers_df.to_sql("customers", con=db_engine, if_exists="append", index=False)
        print(f"  Loaded {len(customers_df)} customers.")

        # D. Products
        print("- Loading 'products'...")
        products_df = df[[
            "product_card_id", "product_name", "product_category_id", 
            "product_price", "product_status"
        ]].drop_duplicates(subset=["product_card_id"])
        # Ensure category ids exist in categories to prevent foreign key errors
        products_df = products_df[products_df["product_category_id"].isin(categories_df["category_id"])]
        products_df.to_sql("products", con=db_engine, if_exists="append", index=False)
        print(f"  Loaded {len(products_df)} products.")

        # E. Orders
        print("- Loading 'orders'...")
        orders_df = df[[
            "order_id", "order_customer_id", "order_date", "shipping_date", 
            "shipping_mode", "days_for_shipping_real", "days_for_shipping_scheduled", 
            "delivery_status", "late_delivery_risk", "order_status", "type", 
            "market", "order_region", "order_country", "order_state", "order_city", 
            "order_zipcode", "latitude", "longitude"
        ]].drop_duplicates(subset=["order_id"])
        # Ensure customer ids exist in customers
        orders_df = orders_df[orders_df["order_customer_id"].isin(customers_df["customer_id"])]
        orders_df.to_sql("orders", con=db_engine, if_exists="append", index=False)
        print(f"  Loaded {len(orders_df)} orders.")

        # F. Order Items
        print("- Loading 'order_items'...")
        order_items_df = df[[
            "order_item_id", "order_id", "order_item_cardprod_id", 
            "order_item_quantity", "order_item_product_price", "order_item_discount", 
            "order_item_discount_rate", "order_item_total", "sales", 
            "benefit_per_order", "order_item_profit_ratio"
        ]].drop_duplicates(subset=["order_item_id"])
        # Ensure order and product ids exist in respective tables
        order_items_df = order_items_df[
            order_items_df["order_id"].isin(orders_df["order_id"]) & 
            order_items_df["order_item_cardprod_id"].isin(products_df["product_card_id"])
        ]
        
        # Insert in chunks to avoid large packet issues in MySQL
        order_items_df.to_sql("order_items", con=db_engine, if_exists="append", index=False, chunksize=10000)
        print(f"  Loaded {len(order_items_df)} order items.")

        print("\nAll tables loaded successfully!")
        
        # Validate count
        with db_engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(*) FROM orders"))
            count = result.scalar()
            print(f"Verification: Found {count} orders loaded in database.")
            
    except Exception as e:
        print(f"Error loading normalized tables: {e}")

if __name__ == "__main__":
    load_database()