import argparse
import datetime
import io
import json
import os
import re
import sys
import subprocess
import zipfile

# ==========================================
# 🛠️ RUNTIME DEPENDENCY INSTALLATION
# ==========================================
try:
    import geopandas as gpd
except ImportError:
    print("📦 Installing geopandas dynamically...", flush=True)
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", 
        "geopandas==0.10.2", "shapely<2.0.0", "fiona<1.9.0"
    ])
    import geopandas as gpd

import pandas as pd
import requests

DEFAULT_OUTPUT_DIR = "raw_data"

# ==========================================
# 🔑 ENVIRONMENT API KEYS
# ==========================================
# Keys are read directly from CI/CD Environment Variables
DATA_GOV_API_KEY = os.getenv("DATA_GOV_API_KEY", "")
DATAMALL_API_KEY = os.getenv("DATAMALL_API_KEY", "")
ONEMAP_TOKEN = os.getenv("ONEMAP_TOKEN", "")
ONEMAP_EMAIL = os.getenv("ONEMAP_EMAIL", "")
ONEMAP_PASSWORD = os.getenv("ONEMAP_PASSWORD", "")

# Dynamic Headers
DATA_GOV_HEADERS = {}
DATAMALL_HEADERS = {}
ONEMAP_HEADERS = {}

# ==========================================
# 📦 DATASETS CONFIGURATION
# ==========================================
DATA_GOV_DATASETS = [
    {
        "id": "d_8594ae9ff96d0c708bc2af633048edfb",
        "name": "subzone_boundary_mp19",
        "file": "subzone_boundary_mp19.geojson",
    },
    {
        "id": "d_d95ae740c0f8961a0b10435836660ce0",
        "name": "pop_age_sex_census2020",
        "file": "pop_age_sex_census2020.csv",
    },
    {
        "id": "d_7f243956483d5901f237e6f87b096636",
        "name": "pop_dwelling_census2020",
        "file": "pop_dwelling_census2020.csv",
    },
    {
        "id": "d_9a6bdc9d93bd041eb0cfbb6a8cb3248f",
        "name": "rail_stations_mp19",
        "file": "rail_stations_mp19.geojson",
    },
    {
        "id": "d_90d86daa5bfaa371668b84fa5f01424f",
        "name": "land_use_mp19",
        "file": "land_use_mp19.geojson",
    },
    {
        "id": "d_688b934f82c1059ed0a6993d2a829089",
        "name": "general_info_schools",
        "file": "general_info_schools.csv",
    },
]

ONEMAP_THEMES = [
    "ssot_hawkercentres",
    "museum",
    "votg_phpc",
    "libraries",
    "communityclubs",
    "monuments",
    "historicsites",
    "hotels",
    "tourism",
    "nationalparks",
]

# ==========================================
# 🔐 INITIALIZE HEADERS & AUTH
# ==========================================
def initialize_headers():
    global DATA_GOV_HEADERS, DATAMALL_HEADERS, ONEMAP_HEADERS
    
    token = ONEMAP_TOKEN

    # Generate fresh JWT Token if credentials exist
    if ONEMAP_EMAIL and ONEMAP_PASSWORD:
        print("🔐 Authenticating with OneMap API...")
        try:
            auth_res = requests.post(
                "https://www.onemap.gov.sg/api/auth/post/getToken",
                json={"email": ONEMAP_EMAIL, "password": ONEMAP_PASSWORD},
                headers={"Content-Type": "application/json"},
                timeout=15
            )
            if auth_res.status_code == 200:
                token = auth_res.json().get("access_token")
                print("✅ Fresh OneMap JWT Token generated successfully!")
        except Exception as err:
            print(f"⚠️ OneMap Auth failed ({err}). Falling back to ONEMAP_TOKEN variable.")

    clean_datamall_key = re.sub(r"[^\x00-\x7F]+", "", str(DATAMALL_API_KEY)).strip()
    clean_onemap_token = str(token).replace("Bearer", "").strip()

    DATA_GOV_HEADERS = {"x-api-key": DATA_GOV_API_KEY, "Content-Type": "application/json"}
    DATAMALL_HEADERS = {"AccountKey": clean_datamall_key, "accept": "application/json"}
    ONEMAP_HEADERS = {"Authorization": f"Bearer {clean_onemap_token}", "accept": "application/json"}


# ==========================================
# 🛠️ HELPER FUNCTIONS FOR API FETCHING
# ==========================================
def fetch_poll_download_dataset(dataset_id, filepath):
    poll_url = f"https://api-open.data.gov.sg/v1/public/api/datasets/{dataset_id}/poll-download"
    try:
        response = requests.get(poll_url, headers=DATA_GOV_HEADERS, timeout=15)
        json_data = response.json()
        if json_data.get("code") != 0:
            print(f"❌ Error polling dataset {dataset_id}: {json_data.get('errMsg')}")
            return False

        download_url = json_data["data"]["url"]
        data_response = requests.get(download_url, timeout=30)

        with open(filepath, "w", encoding="utf-8") as f:
            f.write(data_response.text)
        print(f"✅ Saved to {filepath}")
        return True
    except Exception as e:
        print(f"❌ Failed to process dataset {dataset_id}: {e}")
        return False


def fetch_lta_od_csv(endpoint_type, filepath):
    url = f"https://datamall2.mytransport.sg/ltaodataservice/{endpoint_type}"
    try:
        response = requests.get(url, headers=DATAMALL_HEADERS, timeout=15)
        if response.status_code != 200:
            return None

        download_url = response.json()["value"][0]["Link"]
        zip_res = requests.get(download_url, timeout=60)

        if zip_res.status_code == 200:
            with zipfile.ZipFile(io.BytesIO(zip_res.content)) as z:
                with z.open(z.namelist()[0]) as f:
                    df = pd.read_csv(f)
                    df.to_csv(filepath, index=False)
                    print(f"✅ Saved {len(df):,} records to: {filepath}")
                    return len(df)
        return None
    except Exception as e:
        print(f"❌ Error fetching {endpoint_type}: {e}")
        return None


def fetch_lta_bus_stops(filepath):
    url = "https://datamall2.mytransport.sg/ltaodataservice/BusStops"
    all_stops, skip = [], 0
    try:
        while True:
            res = requests.get(f"{url}?$skip={skip}", headers=DATAMALL_HEADERS, timeout=15)
            if res.status_code != 200:
                break
            data = res.json().get("value", [])
            if not data:
                break
            all_stops.extend(data)
            skip += 500

        if all_stops:
            df = pd.DataFrame(all_stops)
            df.to_csv(filepath, index=False)
            print(f"✅ Saved {len(df):,} bus stop locations to: {filepath}")
            return len(df)
        return None
    except Exception as e:
        print(f"❌ Error fetching bus stops: {e}")
        return None


def fetch_onemap_theme(theme_name, filepath):
    url = f"https://www.onemap.gov.sg/api/public/themesvc/retrieveTheme?queryName={theme_name}"
    try:
        res = requests.get(url, headers=ONEMAP_HEADERS, timeout=15)
        if res.status_code == 200:
            data = res.json().get("SrchResults", [])
            pois = data[1:] if len(data) > 1 else []
            if pois:
                df = pd.DataFrame(pois)
                df.to_csv(filepath, index=False)
                print(f"✅ Saved {len(pois):,} POIs to: {filepath}")
                return len(pois)
        return None
    except Exception as e:
        print(f"❌ Exception fetching theme '{theme_name}': {e}")
        return None


# ==========================================
# 🚀 MAIN PIPELINE EXECUTION
# ==========================================
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=str, default=DEFAULT_OUTPUT_DIR, help="Local directory to store fetched raw data")
    args, _ = parser.parse_known_args()

    print(f"🚀 Starting Local Raw Data Fetch -> Output Directory: '{args.output_dir}'")

    initialize_headers()

    raw_data_dir = args.output_dir
    os.makedirs(raw_data_dir, exist_ok=True)

    # 1. Data.gov.sg Batch Fetch
    print("\n📦 Fetching Data.gov.sg Datasets...")
    for item in DATA_GOV_DATASETS:
        out_file = os.path.join(raw_data_dir, item["file"])
        fetch_poll_download_dataset(item["id"], out_file)

    # 2. LTA DataMall Endpoints
    print("\n📦 Fetching LTA DataMall Endpoints...")
    fetch_lta_od_csv("PV/ODBus", os.path.join(raw_data_dir, "lta_od_bus.csv"))
    fetch_lta_od_csv("PV/ODTrain", os.path.join(raw_data_dir, "lta_od_train.csv"))
    fetch_lta_bus_stops(os.path.join(raw_data_dir, "lta_bus_stops.csv"))

    # 3. OneMap POI Themes
    print(f"\n📦 Fetching {len(ONEMAP_THEMES)} OneMap POI Themes...")
    for theme_api_name in ONEMAP_THEMES:
        out_file = os.path.join(raw_data_dir, f"onemap_{theme_api_name}.csv")
        fetch_onemap_theme(theme_api_name, out_file)

    print(f"\n🎉 Local Data Ingestion Complete! All raw data stored in '{raw_data_dir}'")