import os, io, time, requests
from datetime import datetime
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed

DEFAULT_LOCAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'output', 'pincode_master.parquet'))
DEFAULT_S3_KEY = 'pincode_master/pincode_master.parquet'


def clean_pincode(val):
    """Normalize input into a 6-digit Indian postal code string, or return None."""
    if pd.isna(val):
        return None
    val_str = str(val).split('.')[0].strip()
    return val_str if (len(val_str) == 6 and val_str.isdigit()) else None


def _empty_result(pincode, status):
    """Return a blank geocode result row for failed or unresolved pincodes."""
    return {
        'pincode': str(pincode), 'latitude': None, 'longitude': None,
        'formatted_address': None, 'locality': None, 'district': None,
        'state': None, 'status': status,
        'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }


def batch_geocode_pincodes(pincodes, api_key, max_workers=10):
    """
    Geocode a list of pincodes concurrently using the Google Maps Geocoding API.
    Uses a ThreadPoolExecutor with up to 10 workers, staying within Google QPS limits.
    Returns a DataFrame with one row per pincode.
    """
    def geocode_one(pincode, session):
        """Inner function: make a single geocode API call and parse the result."""
        try:
            resp = session.get(
                "https://maps.googleapis.com/maps/api/geocode/json",
                params={'address': f"{pincode}, India", 'key': api_key},
                timeout=10
            ).json()
        except Exception as e:
            return _empty_result(pincode, f"Error: {e}")

        status = resp.get('status')
        if status == 'OK' and resp.get('results'):
            result = resp['results'][0]
            lat = result['geometry']['location']['lat']
            lng = result['geometry']['location']['lng']
            components = {c['types'][0]: c['long_name'] for c in result.get('address_components', [])}
            return {
                'pincode': str(pincode),
                'latitude': round(float(lat), 6),
                'longitude': round(float(lng), 6),
                'formatted_address': result.get('formatted_address'),
                'locality': components.get('sublocality') or components.get('locality') or components.get('administrative_area_level_2'),
                'district': components.get('administrative_area_level_2'),
                'state': components.get('administrative_area_level_1'),
                'status': 'OK',
                'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
        return _empty_result(pincode, status)

    total = len(pincodes)
    results = []
    print(f"Geocoding {total} unique pincodes via Google Maps API (workers={max_workers})...")
    start_time = time.time()

    with requests.Session() as session:
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(geocode_one, pin, session): pin for pin in pincodes}
            for i, future in enumerate(as_completed(futures), 1):
                results.append(future.result())
                if i % 500 == 0 or i == total:
                    elapsed = time.time() - start_time
                    print(f"  Geocoded {i}/{total} ({i / elapsed:.1f} pins/sec)...")

    print(f"Completed geocoding in {time.time() - start_time:.2f} seconds.")
    return pd.DataFrame(results)


def _load_existing_master(local_path, s3_client, bucket_name, s3_key):
    """
    Load pincode_master.parquet from local disk.
    Falls back to downloading from S3 if the local file is absent or unreadable.
    Returns a DataFrame, or None if no existing master is found.
    """
    if os.path.exists(local_path):
        try:
            df = pd.read_parquet(local_path)
            print(f"Loaded existing pincode master ({len(df)} records) from local disk.")
            return df
        except Exception as e:
            print(f"Warning: Could not read local parquet ({e}). Checking S3...")

    if s3_client and bucket_name:
        try:
            buf = io.BytesIO()
            s3_client.download_fileobj(bucket_name, s3_key, buf)
            buf.seek(0)
            df = pd.read_parquet(buf)
            print(f"Downloaded existing pincode master ({len(df)} records) from S3.")
            return df
        except Exception as e:
            print(f"No existing master found in S3 (first run or missing): {e}")

    return None


def _save_and_upload(df, local_path, s3_client, bucket_name, s3_key):
    """
    Save the updated pincode master locally, upload it to S3,
    then remove the local file to avoid stale copies on disk.
    """
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    df.to_parquet(local_path, index=False, engine='pyarrow')
    print(f"Saved updated pincode master ({len(df)} records) locally.")

    if s3_client and bucket_name:
        buf = io.BytesIO()
        df.to_parquet(buf, index=False, engine='pyarrow')
        buf.seek(0)
        print(f"Uploading pincode master to s3://{bucket_name}/{s3_key} ...")
        s3_client.upload_fileobj(buf, bucket_name, s3_key)
        print("Successfully uploaded pincode_master to S3.")

        os.remove(local_path)
        print("Cleaned up local pincode_master.parquet after S3 upload.")


def sync_pincode_master(raw_pincodes, s3_client=None, bucket_name=None, api_key=None,
                        local_path=DEFAULT_LOCAL_PATH, s3_key=DEFAULT_S3_KEY):
    """
    Incrementally sync the pincode master parquet with any new pincodes from the current run.

    Workflow:
      1. Load existing pincode_master.parquet from local disk or S3.
      2. Identify pincodes in this run that are not yet in the master.
      3. Geocode only the new pincodes via Google Maps API.
      4. Append, save locally, upload to S3, and remove the local copy.

    Returns the full up-to-date pincode master as a DataFrame.
    """
    from dotenv import load_dotenv
    load_dotenv()

    if not api_key:
        api_key = os.getenv('GOOGLE_MAPS_API_KEY')
    if not api_key:
        raise ValueError("GOOGLE_MAPS_API_KEY is missing. Set it in .env or pass api_key.")

    # Step 1: Load existing master
    existing_df = _load_existing_master(local_path, s3_client, bucket_name, s3_key)
    existing_pins = set()
    if existing_df is not None and not existing_df.empty:
        existing_df['pincode'] = existing_df['pincode'].astype(str).str.strip()
        existing_pins = set(existing_df['pincode'].unique())

    # Step 2: Identify new pincodes
    clean_pins = {cp for p in raw_pincodes if (cp := clean_pincode(p))}
    unresolved_pins = list(clean_pins - existing_pins)
    print(f"Pincode Sync: {len(clean_pins)} in run | {len(existing_pins)} in master | {len(unresolved_pins)} new to geocode.")

    # Step 3: Geocode and update if new pincodes exist
    if not unresolved_pins:
        print("All pincodes already exist in master. Zero API calls made.")
        return existing_df

    new_df = batch_geocode_pincodes(unresolved_pins, api_key)
    updated_df = pd.concat([existing_df, new_df], ignore_index=True) if existing_df is not None and not existing_df.empty else new_df
    updated_df = updated_df.drop_duplicates(subset=['pincode'], keep='last')

    # Step 4: Persist
    _save_and_upload(updated_df, local_path, s3_client, bucket_name, s3_key)
    return updated_df