gemini.py

# gemini.py
import os
import re
import io
import json
import datetime
import pandas as pd
from PIL import Image
import google.generativeai as genai
from dateutil.parser import parse
import argparse
import csv
import uuid

# --- Local Imports ---
from supabase_client import get_supabase_client, fetch_races_db_fields
from config import GEMINI_API_KEY, OUTPUT_DIR, GEMINI_ANALYTICS_DIR

# --- Configuration ---
DEFAULT_INPUT_FOLDER = 'gemini_input_temp'
OUTPUT_CSV_FILENAME_PREFIX = 'Gemini_Image_Extraction'
CACHE_FILE = 'processed_images.log'
ALLOWED_EXTENSIONS = ('.png', '.jpeg', '.jpg')
CSV_HEADERS =['event', 'festivalName', 'imageURL', 'raceVideo', 'type', 'date', 'city', 'organiser','participationType', 'firstEdition', 'lastEdition', 'countEditions', 'mode', 'raceAccredition','theme', 'numberOfparticipants', 'startTime', 'scenic', 'registrationCost', 'ageLimitation','eventWebsite', 'organiserWebsite', 'bookingLink', 'newsCoverage', 'lastDate','participationCriteria', 'refundPolicy', 'swimDistance', 'swimType', 'swimmingLocation','waterTemperature', 'swimCoursetype', 'swimCutoff', 'swimRoutemap', 'cyclingDistance','cyclingElevation', 'cyclingSurface', 'cyclingElevationgain', 'cycleCoursetype', 'cycleCutoff','cyclingRoutemap', 'runningDistance', 'runningElevation', 'runningSurface', 'runningElevationgain','runningElevationloss', 'runningCoursetype', 'runCutoff', 'runRoutemap', 'organiserRating','triathlonType', 'standardTag', 'region', 'approvalStatus', 'difficultyLevel', 'month','primaryKey', 'latitude', 'longitude', 'country', 'editionYear', 'aidStations','restrictedTraffic', 'user_id', 'femaleParticpation', 'jellyFishRelated','registrationOpentag', 'eventConcludedtag', 'state', 'nextEdition']
CHOICE_FIELDS = {"participationType":["Individual", "Relay", "Group"], "mode": ["Virtual", "On-Ground"], "runningSurface":["Road", "Trail", "Track", "Road + Trail"], "runningCourseType":["Single Loop", "Multiple Loop", "Out and Back", "Point to Point"], "region":["West India", "Central and East India", "North India", "South India", "Nepal", "Bhutan", "Sri Lanka"], "runningElevation":["Flat", "Rolling", "Hilly", "Skyrunning"], "type":["Triathlon", "Aquabike", "Aquathlon", "Duathlon", "Run", "Cycling", "Swimathon"], "swimType":["Lake", "Beach", "River", "Pool"], "swimCoursetype":["Single Loop", "Multiple Loops", "Out and Back", "Point to Point"], "cyclingElevation":["Flat", "Rolling", "Hilly"], "cycleCoursetype":["Single Loop", "Multiple Loops", "Out and Back", "Point to Point"], "triathlonType":["Super Sprint", "Sprint Distance", "Olympic Distance", "Half Iron(70.3)", "Iron Distance (140.6)","Ultra Distance"], "standardTag":["Standard", "Non Standard"], "restrictedTraffic":["Yes", "No"], "jellyFishRelated": ["Yes", "No"], "approvalStatus":["Approved", "Pending Approval"]}

# --- Output Persistence ---
def save_output_to_supabase(filepath: str, agent_mode: str, event_type: str = "Image Extraction"):
    supabase = get_supabase_client()
    if not supabase: return

    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            run_data = list(reader)
        
        if not run_data: return

        payload = {
            "agent_mode": agent_mode,
            "event_type": event_type,
            "filename": os.path.basename(filepath),
            "run_data": run_data,
            "event_count": len(run_data)
        }
        
        supabase.table('run_outputs').insert(payload).execute()
        print("[SUCCESS] Successfully saved run output to Supabase.")

    except Exception as e:
        print(f"[ERROR] Failed to save output file to Supabase: {e}")

# --- Data Cleaning and Formatting Helpers ---
def clean_value(value):
    if value is None or str(value).strip().upper() in["NA", "N/A", "NONE", "NOT SPECIFIED", ""]: return ""
    return str(value).strip()

def validate_choice(value, options):
    cleaned_value = clean_value(value)
    if not cleaned_value: return ""
    for option in options:
        if cleaned_value.lower() == option.lower(): return option
    return ""

def format_date_value(date_str):
    cleaned_str = clean_value(date_str)
    if not cleaned_str: return ""
    try:
        dt = parse(cleaned_str, fuzzy=True, dayfirst=True)
        if dt and (dt.year >= datetime.datetime.now().year - 1):
            return dt.strftime("%d-%m-%Y")
        return ""
    except (ValueError, TypeError): return ""

def format_time_value(time_str):
    cleaned_str = clean_value(time_str)
    if not cleaned_str: return ""
    try:
        dt = parse(cleaned_str, fuzzy=True)
        return dt.strftime("%H:%M")
    except (ValueError, TypeError): return ""

def extract_numeric(value_str):
    cleaned_str = clean_value(str(value_str))
    if not cleaned_str: return ""
    match = re.search(r'(\d+\.?\d*|\.\d+)', cleaned_str)
    return match.group(1) if match else ""

def extract_registration_cost(value_str):
    cleaned_str = clean_value(value_str)
    if not cleaned_str: return ""
    if "free" in cleaned_str.lower(): return "0"
    cleaned_str = cleaned_str.replace(',', '')
    match = re.search(r'(\d+)', cleaned_str)
    return match.group(1) if match else ""

# --- Core Functions ---
def get_gemini_response(image_path: str, filename: str) -> list:
    if not GEMINI_API_KEY: raise ValueError("GEMINI_API_KEY not found in configuration.")

    try:
        genai.configure(api_key=GEMINI_API_KEY)
        with Image.open(image_path) as img:
            img_byte_arr = io.BytesIO()
            if img.mode == 'RGBA': img = img.convert('RGB')
            img.save(img_byte_arr, format='JPEG')
            image_bytes = img_byte_arr.getvalue()

        image_part = {"mime_type": "image/jpeg", "data": image_bytes}
        model = genai.GenerativeModel('gemini-1.5-flash-latest')

        # --- DYNAMIC CONTEXT INJECTION ---
        db_fields = fetch_races_db_fields()
        meta_map = {row['field']: row for row in db_fields} if db_fields else {}
        
        gemini_fields =['event', 'type', 'date', 'city', 'organiser', 'registrationCost', 'lastDate', 'startTime', 'eventWebsite', 'runningDistance', 'cyclingDistance', 'swimDistance']
        field_instructions =[]
        for f in gemini_fields:
            meta = meta_map.get(f)
            if meta:
                inst = f"- {f}: ({meta.get('display_name', f)})"
                if meta.get('question_text'): inst += f" Context: {meta['question_text']}."
                opts = meta.get('data_options')
                if opts:
                    opt_str = opts.replace('\n', ', ') if isinstance(opts, str) else str(opts)
                    inst += f" Options: [{opt_str}]."
                elif f in CHOICE_FIELDS:
                    inst += f" Options: [{', '.join(CHOICE_FIELDS[f])}]."
                dt = meta.get('data_type')
                df = meta.get('data_format')
                if dt or df: inst += f" Format: {dt or ''} {df or ''}."
                field_instructions.append(inst)
            else:
                field_instructions.append(f"- {f}: Extract the {f}.")

        fields_block = "\n        ".join(field_instructions)

        prompt = f"""
        You are a highly intelligent data extraction assistant specializing in athletic events. Analyze the provided image of an event poster or banner and extract the following information.
        
        CRITICAL REQUIREMENT:
        You MUST return a JSON LIST of objects. IF the event has multiple race distances or variants (e.g., 5k, 10k, 21.1k), you MUST create a SEPARATE JSON object in the list for EACH distance. 
        Append the distance to the 'event' name for that specific row (e.g., "City Run 2025 - 5k").
        
        Fields to Extract for EACH object (Strictly adhere to definitions):
        {fields_block}
        
        Respond ONLY with a valid JSON LIST. Do not include any text or markdown formatting before or after the JSON.
        
        Example Response for an event with 2 distances:[
          {{
            "event": "City Runners Marathon 2025 - 10k",
            "type": "Run",
            "date": "23-11-2025",
            "city": "Mumbai",
            "registrationCost": "1500",
            "runningDistance": "10"
          }},
          {{
            "event": "City Runners Marathon 2025 - 21.1k",
            "type": "Run",
            "date": "23-11-2025",
            "city": "Mumbai",
            "registrationCost": "2000",
            "runningDistance": "21.1"
          }}
        ]
        """

        response = model.generate_content([prompt, image_part], request_options={"timeout": 120})
        raw_text = response.text.strip()
        
        os.makedirs(GEMINI_ANALYTICS_DIR, exist_ok=True)
        analytics_file = os.path.join(GEMINI_ANALYTICS_DIR, f"gemini_image_{filename}_{uuid.uuid4().hex[:6]}.json")
        with open(analytics_file, 'w', encoding='utf-8') as f:
            f.write(raw_text)

        clean_response_text = re.sub(r'^```json\s*|\s*```$', '', raw_text, flags=re.MULTILINE)
        data = json.loads(clean_response_text)
        
        if isinstance(data, dict):
            return [data]
        return data

    except Exception as e:
        print(f"  ->[ERROR] Error calling Gemini API for {os.path.basename(image_path)}: {e}")
        return[]

def process_image_data(raw_data: dict) -> dict:
    processed_row = {header: "" for header in CSV_HEADERS}

    for key, value in raw_data.items():
        if key in processed_row:
            if key in CHOICE_FIELDS:
                processed_row[key] = validate_choice(value, CHOICE_FIELDS[key])
            elif "Date" in key or "date" in key:
                processed_row[key] = format_date_value(value)
            elif "Time" in key:
                processed_row[key] = format_time_value(value)
            elif "Cost" in key:
                processed_row[key] = extract_registration_cost(value)
            elif "Distance" in key:
                if isinstance(value, list): processed_row[key] = ", ".join(map(str, value))
                else: processed_row[key] = extract_numeric(value)
            else:
                processed_row[key] = clean_value(value)

    if processed_row.get("date"):
        try:
            dt = parse(processed_row["date"], fuzzy=True, dayfirst=True)
            processed_row["month"] = dt.strftime("%B")
            processed_row["editionYear"] = dt.strftime("%Y")
        except (ValueError, TypeError): pass

    return processed_row

def main(output_dir_override=None, input_dir_override=None, output_filename=None):
    print("--- Gemini Image Processor Script ---")
    effective_input_dir = input_dir_override if input_dir_override else DEFAULT_INPUT_FOLDER
    print(f"[INFO] Reading images from: {os.path.abspath(effective_input_dir)}")

    output_dir = output_dir_override or OUTPUT_DIR
    if not os.path.exists(output_dir): os.makedirs(output_dir)

    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    base_name = output_filename if output_filename else f"{OUTPUT_CSV_FILENAME_PREFIX}_{timestamp}"
    final_output_path = os.path.join(output_dir, f"{base_name}.csv")

    if not os.path.exists(effective_input_dir):
        print(f"[ERROR] Input directory '{effective_input_dir}' not found. Cannot proceed.")
        return

    try:
        with open(CACHE_FILE, 'r', encoding='utf-8') as f: processed_images = set(f.read().splitlines())
    except FileNotFoundError: processed_images = set()

    all_images =[f for f in os.listdir(effective_input_dir) if f.lower().endswith(ALLOWED_EXTENSIONS)]
    new_images =[f for f in all_images if f not in processed_images]

    if not new_images:
        print("[INFO] No new images to process. Exiting.")
        return

    print(f"[INFO] Found {len(new_images)} new image(s) to process.")
    new_data_rows = []
    processed_this_run =[]

    for image_name in new_images:
        image_path = os.path.join(effective_input_dir, image_name)
        print(f"\n[INFO] Processing '{image_name}'...")
        
        raw_data_list = get_gemini_response(image_path, image_name)

        if not raw_data_list:
            print(f"  -> [WARNING] Skipping {image_name} due to API error or empty response.")
            continue

        for raw_data in raw_data_list:
            processed_row = process_image_data(raw_data)
            new_data_rows.append(processed_row)
            
        processed_this_run.append(image_name)
        print(f"  ->[SUCCESS] Extracted {len(raw_data_list)} variants for '{image_name}'.")

    if not new_data_rows:
        print("\n[INFO] No data was successfully extracted from any new images. Exiting.")
        return

    df = pd.DataFrame(new_data_rows)
    df = df[CSV_HEADERS]

    df.to_csv(final_output_path, mode='w', header=True, index=False, encoding='utf-8')
    print(f"\n[SUCCESS] Data for {len(new_data_rows)} variant(s) successfully saved to '{final_output_path}'.")

    save_output_to_supabase(final_output_path, agent_mode="gemini")

    with open(CACHE_FILE, 'a', encoding='utf-8') as f:
        for image_name in processed_this_run: f.write(f"{image_name}\n")

    print("\n--- Script Finished ---")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Crawl4AI Gemini Image Processor")
    parser.add_argument("--output-dir", type=str, default=OUTPUT_DIR, help="Directory to save output files.")
    parser.add_argument("--input-dir", type=str, required=True, help="Directory containing input images.")
    parser.add_argument("--output-filename", type=str, help="Custom name for the output file (without extension).")
    args = parser.parse_args()

    main(output_dir_override=args.output_dir, input_dir_override=args.input_dir, output_filename=args.output_filename)