import requests
# import polyline # Not used in the provided snippet
import json
import psycopg2
import psycopg2.extensions # Not strictly necessary for this code, but fine to keep
from datetime import datetime, timezone # timezone not used, but datetime is
# from geojson import Point, Feature, FeatureCollection, dump # Not used
# import pandas as pd # Not used
from requests.packages.urllib3.exceptions import InsecureRequestWarning

# --- Configuration ---
DB_HOST = 'localhost'
DB_NAME = 'nws'
DB_USER = 'nws'
DB_PASSWORD = 'nws' # Consider using environment variables or a secrets manager

# API endpoint (contains specific coordinates, ensure this is intended)
OHGO_API_URL = "https://api.ohgo.com/roadmarkers/TrafficSpeedAndAlertMarkers?pointData={%22lowLongitude%22:-83.55941510027988,%22highLongitude%22:-80.38711285418613,%22lowLatitude%22:38.243663436655325,%22highLatitude%22:40.169517342197835,%22routeDirection%22:%22%22,%22routeName%22:%22%22}"

# Optional Proxies
PROXIES = {"http": "http://nws:nws@localhost:9000",
           "https": "http://nws:nws@localhost:9000"} # Add https if needed

# Disable insecure request warnings (use cautiously)
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

# --- Main Logic ---
def fetch_ohgo_data(session, url):
    """Fetches data from the OHGO API."""
    print(f"INFO: Fetching data from {url}...")
    try:
        response = session.get(url, timeout=30) # Add a timeout
        response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx)
        print("INFO: Data fetched successfully.")
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"ERROR: Failed to fetch data from API: {e}")
        return None
    except json.JSONDecodeError as e:
        print(f"ERROR: Failed to parse JSON response: {e}")
        print(f"Response text: {response.text[:500]}...") # Log part of the response
        return None

def prepare_data_for_upsert(ohgo_data):
    """Processes API data and prepares it for database insertion."""
    ohgoclosures_for_upsert = []
    # Get timestamp ONCE for this batch, use Python datetime object
    current_timestamp_obj = datetime.utcnow()

    if not ohgo_data or 'TrafficAlertMarkers' not in ohgo_data:
        print("WARNING: No 'TrafficAlertMarkers' found in API response.")
        return []

    alerts = ohgo_data['TrafficAlertMarkers']
    print(f"INFO: Processing {len(alerts)} alerts from API.")

    for j in alerts:
        category = j.get('Category')
        # Filter for relevant categories
        if category == "Flooding" or category == "Weather":
            # Prepare tuple with exactly 10 items for the efficient UPSERT SQL
            info_tuple = (
                j.get('Latitude'),
                j.get('Longitude'),
                j.get('Id'),
                j.get('Location'),
                j.get('Description'),
                category, # Use variable already fetched
                j.get('RoadStatus'),
                j.get('RouteName'),
                current_timestamp_obj, # Use datetime object for 'start'
                current_timestamp_obj  # Use datetime object for 'lastupdate'
            )
            # Basic validation: check if critical fields like ID are present
            if info_tuple[2] is None:
                print(f"WARNING: Skipping record due to missing ID: {j}")
                continue
            if info_tuple[0] is None or info_tuple[1] is None:
                 print(f"WARNING: Skipping record {info_tuple[2]} due to missing Lat/Lon: {j}")
                 continue

            ohgoclosures_for_upsert.append(info_tuple)

    print(f"INFO: Prepared {len(ohgoclosures_for_upsert)} records for database.")
    return ohgoclosures_for_upsert

def update_database(conn, data_to_upsert):
    """Performs batch upsert and subsequent updates in the database."""
    if not data_to_upsert:
        print("INFO: No data to update in the database.")
        return

    # SQL for batch UPSERT using ON CONFLICT...DO UPDATE SET...EXCLUDED
    upsert_sql = """
        INSERT INTO ohgo (
            lat, lon, id, location, description, category,
            roadstatus, routename, start, lastupdate
        )
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) -- 10 placeholders
        ON CONFLICT (id) DO UPDATE SET
            lat = EXCLUDED.lat,
            lon = EXCLUDED.lon,
            location = EXCLUDED.location,
            description = EXCLUDED.description,
            category = EXCLUDED.category,
            roadstatus = EXCLUDED.roadstatus,
            lastupdate = EXCLUDED.lastupdate;
    """

    # Use context managers for cursor and transaction handling
    try:
        with conn: # Start transaction
            with conn.cursor() as cursor:
                # --- 1. Batch Upsert ---
                print(f"INFO: Upserting {len(data_to_upsert)} records...")
                cursor.executemany(upsert_sql, data_to_upsert)
                # Note: rowcount might be unreliable with executemany depending on driver/DB version
                print(f"INFO: Upsert complete. {cursor.rowcount} rows affected (Note: count may be inaccurate).")

                # --- 2. Post-Processing Updates ---
                print("INFO: Performing post-processing updates...")

                # Update geometry (Requires spatial index on ohgo.geom for performance)
                print("INFO: Updating geometry...")
                cursor.execute("""
                    UPDATE ohgo
                    SET geom = ST_SetSRID(ST_MakePoint(lon, lat), 4326)
                    WHERE lat IS NOT NULL AND lon IS NOT NULL AND geom IS NULL;
                """)
                print(f"INFO: Updated geometry for {cursor.rowcount} records.")

                # Update endtime based on lastupdate
                print("INFO: Updating end times...")
                cursor.execute("""
                    UPDATE ohgo
                    SET endtime = NULL
                    WHERE lastupdate > now() - interval '0.25 hours';
                """)
                print(f"INFO: Reset endtime for {cursor.rowcount} recent records.")
                cursor.execute("""
                    UPDATE ohgo
                    SET endtime = lastupdate
                    WHERE endtime IS NULL AND lastupdate < now() - interval '1 hours';
                """)
                print(f"INFO: Set endtime for {cursor.rowcount} older records.")

                # Update administrative boundaries (CRITICAL: Requires spatial indexes!)
                print("INFO: Updating County/CWA/State...")
                # Add WHERE geom IS NOT NULL for efficiency
                cursor.execute("""
                    UPDATE ohgo o SET county = c.countyname
                    FROM public.county c
                    WHERE o.county IS NULL AND o.geom IS NOT NULL
                      AND ST_Contains(c.geom, o.geom);
                """)
                print(f"INFO: Updated county for {cursor.rowcount} records.")
                cursor.execute("""
                    UPDATE ohgo o SET cwa = f.cwa
                    FROM public.fzone f
                    WHERE o.cwa IS NULL AND o.geom IS NOT NULL
                      AND ST_Contains(f.geom, o.geom);
                """)
                print(f"INFO: Updated cwa for {cursor.rowcount} records.")
                cursor.execute("""
                    UPDATE ohgo o SET state = c.state
                    FROM public.county c
                    WHERE o.state IS NULL AND o.geom IS NOT NULL
                      AND ST_Contains(c.geom, o.geom);
                """)
                print(f"INFO: Updated state for {cursor.rowcount} records.")

                # Set default boolean flags (Combined)
                print("INFO: Setting default flags...")
                cursor.execute("""
                    UPDATE ohgo
                    SET lsr = false, hide = false
                    WHERE lsr IS NULL;
                """)
                print(f"INFO: Set default flags for {cursor.rowcount} records.")

                # Clean description (Combined)
                print("INFO: Cleaning descriptions...")
                cursor.execute("""
                    UPDATE ohgo
                    SET description = TRIM(REPLACE(REPLACE(description, 'Use alternate route.', ''), E'\n', ' '))
                    WHERE description LIKE '%Use alternate route.%' OR description LIKE E'%\n%';
                """)
                print(f"INFO: Cleaned description for {cursor.rowcount} records.")

        print("INFO: Database operations successful. Transaction committed.") # Commit happens automatically on exiting 'with conn:'

    except psycopg2.Error as e:
        print(f"ERROR: Database error occurred: {e}")
        print("ERROR: Transaction automatically rolled back.")
        # Optionally re-raise the exception if calling code needs to handle it
        # raise
    except Exception as e:
        print(f"ERROR: An unexpected error occurred during database operations: {e}")
        # Transaction is still rolled back by 'with conn:' exiting on exception
        # raise

# --- Script Execution ---
if __name__ == "__main__":
    conn = None # Initialize connection variable
    try:
        # Establish database connection
        print(f"INFO: Connecting to database '{DB_NAME}' on '{DB_HOST}'...")
        conn = psycopg2.connect(host=DB_HOST, database=DB_NAME, user=DB_USER, password=DB_PASSWORD)
        print("INFO: Database connection successful.")

        # Create a requests session
        S = requests.Session()
        S.verify = False # Disables SSL verification - use with caution
        #S.proxies = PROXIES # Use the defined proxies

        # Fetch data
        ohgo_json_data = fetch_ohgo_data(S, OHGO_API_URL)

        if ohgo_json_data:
            # Prepare data
            prepared_data = prepare_data_for_upsert(ohgo_json_data)

            # Update database
            update_database(conn, prepared_data)

    except psycopg2.OperationalError as e:
        print(f"FATAL: Could not connect to database: {e}")
    except Exception as e:
        print(f"FATAL: An unexpected error occurred: {e}")
    finally:
        # Ensure the connection is closed
        if conn:
            conn.close()
            print("INFO: Database connection closed.")

    print("INFO: Script finished.")