#!/usr/bin/env python3
"""Import BTS CSV data from downloaded ZIPs into the flights table.

Reads in 100K-row chunks, drops indexes during load, uses fast SQLite PRAGMAs.
"""
import os
import sys
import zipfile
from pathlib import Path

import pandas as pd
from sqlalchemy import text

sys.path.insert(0, str(Path(__file__).parent.parent))

from core.database import engine, Base
from app.models import Flight

# BTS PREZIP column names (CamelCase) -- only import what we need
COLUMNS_TO_KEEP = [
    "Year", "Month", "FlightDate",
    "Reporting_Airline", "Flight_Number_Reporting_Airline", "Tail_Number",
    "Origin", "Dest", "CRSDepTime",
    "ArrDelay", "DepDelay", "ArrDel15",
    "CarrierDelay", "WeatherDelay", "NASDelay",
    "SecurityDelay", "LateAircraftDelay",
    "Cancelled", "Diverted", "Distance",
]

# BTS column names -> our DB column names
COLUMN_RENAME = {
    "Year": "year",
    "Month": "month",
    "FlightDate": "fl_date",
    "Reporting_Airline": "carrier",
    "Flight_Number_Reporting_Airline": "flight_num",
    "Tail_Number": "tail_number",
    "Origin": "origin",
    "Dest": "dest",
    "CRSDepTime": "crs_dep_time",
    "ArrDelay": "arr_delay",
    "DepDelay": "dep_delay",
    "ArrDel15": "arr_del15",
    "CarrierDelay": "carrier_delay",
    "WeatherDelay": "weather_delay",
    "NASDelay": "nas_delay",
    "SecurityDelay": "security_delay",
    "LateAircraftDelay": "late_aircraft_delay",
    "Cancelled": "cancelled",
    "Diverted": "diverted",
    "Distance": "distance",
}

CHUNK_SIZE = 100_000


def setup_database():
    print("Setting up database (dropping existing tables)...")
    Base.metadata.drop_all(bind=engine)
    Base.metadata.create_all(bind=engine)

    # Drop non-PK indexes for faster bulk import
    with engine.connect() as conn:
        for idx_name in ["idx_origin", "idx_dest", "idx_origin_dest", "idx_year_month"]:
            conn.execute(text(f"DROP INDEX IF EXISTS {idx_name}"))
        conn.commit()
    print("  Tables created (indexes deferred)")


def set_fast_pragmas():
    # crank up SQLite for bulk loading
    with engine.connect() as conn:
        conn.execute(text("PRAGMA synchronous=OFF"))
        conn.execute(text("PRAGMA journal_mode=MEMORY"))
        conn.execute(text("PRAGMA cache_size=-2000000"))
        conn.commit()


def restore_safe_pragmas():
    with engine.connect() as conn:
        conn.execute(text("PRAGMA synchronous=NORMAL"))
        conn.execute(text("PRAGMA journal_mode=WAL"))
        conn.commit()


def import_zip(zip_path: Path) -> int:
    with zipfile.ZipFile(zip_path) as zf:
        csv_names = [n for n in zf.namelist() if n.endswith(".csv")]
        if not csv_names:
            print(f"  WARNING: No CSV found in {zip_path.name}")
            return 0

        csv_name = csv_names[0]
        total_rows = 0

        with zf.open(csv_name) as csv_file:
            for chunk in pd.read_csv(
                csv_file,
                usecols=lambda col: col in COLUMNS_TO_KEEP,
                dtype={"Origin": str, "Dest": str, "Reporting_Airline": str, "Tail_Number": str},
                chunksize=CHUNK_SIZE,
                low_memory=False,
                encoding="latin-1",
            ):
                chunk = chunk.rename(columns=COLUMN_RENAME)

                # Convert fl_date string to proper format for SQLite
                if "fl_date" in chunk.columns:
                    chunk["fl_date"] = pd.to_datetime(chunk["fl_date"]).dt.date

                # Convert float flags to int
                for col in ["arr_del15", "cancelled", "diverted"]:
                    if col in chunk.columns:
                        chunk[col] = chunk[col].fillna(0).astype(int)

                # CRSDepTime comes as float (e.g. 1430.0), convert to nullable int
                if "crs_dep_time" in chunk.columns:
                    chunk["crs_dep_time"] = pd.to_numeric(chunk["crs_dep_time"], errors="coerce")
                    chunk["crs_dep_time"] = chunk["crs_dep_time"].astype("Int64")

                chunk.to_sql(
                    "flights",
                    con=engine,
                    if_exists="append",
                    index=False,
                    method="multi",
                )
                total_rows += len(chunk)
                print(f"\r  {zip_path.name}: {total_rows:>10,} rows", end="", flush=True)

        print()
        return total_rows


def recreate_indexes():
    # build indexes after bulk load -- way faster than indexing during inserts
    print("Creating indexes...")
    with engine.connect() as conn:
        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_origin ON flights(origin)"))
        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_dest ON flights(dest)"))
        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_origin_dest ON flights(origin, dest)"))
        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_year_month ON flights(year, month)"))
        conn.commit()
    print("  Indexes created")


def main():
    raw_dir = Path(os.getenv("BTS_DOWNLOAD_DIR", "./data/raw"))
    zip_files = sorted(raw_dir.glob("*.zip"))

    if not zip_files:
        print(f"No ZIP files found in {raw_dir.resolve()}")
        print("Run download_bts_data.py first.")
        sys.exit(1)

    print(f"Found {len(zip_files)} ZIP files to import")
    print(f"Database: {engine.url}")
    print()

    setup_database()
    set_fast_pragmas()

    grand_total = 0
    for i, zip_path in enumerate(zip_files, 1):
        print(f"[{i}/{len(zip_files)}] ", end="")
        rows = import_zip(zip_path)
        grand_total += rows

    recreate_indexes()
    restore_safe_pragmas()

    print(f"\nImport complete: {grand_total:,} total flights loaded")

    # Verify
    with engine.connect() as conn:
        count = conn.execute(text("SELECT COUNT(*) FROM flights")).scalar()
        origins = conn.execute(text("SELECT COUNT(DISTINCT origin) FROM flights")).scalar()
        dests = conn.execute(text("SELECT COUNT(DISTINCT dest) FROM flights")).scalar()
        print(f"Verification: {count:,} rows, {origins} unique origins, {dests} unique destinations")


if __name__ == "__main__":
    main()
