web_research.py

# web_research.py

import os
import re
import csv
import json
import logging
from datetime import datetime
from urllib.parse import urljoin

from config import (
    MISTRAL_API_KEY, SEARCH_API_KEY, CSE_ID,
    OUTPUT_DIR, KNOWLEDGE_CACHE_DIR, APP_VERSION
)
from schemas import (
    TRIATHLON_SCHEMA, RUNNING_SCHEMA, SWIMMING_SCHEMA, DUATHLON_SCHEMA,
    AQUATHLON_SCHEMA, AQUABIKE_SCHEMA, CYCLING_SCHEMA, FITNESS_RACING_SCHEMA
)
from agent import MistralAnalystAgent, AgentInitializationError
from utils import (
    serialize_knowledge_base, deserialize_knowledge_base, format_final_row, SafeCSVWriter
)
from supabase_client import get_supabase_client

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def save_output_to_supabase(filepath: str, agent_mode: str, event_type: str = None):
    supabase = get_supabase_client()
    if not supabase: return

    try:
        if not os.path.exists(filepath):
            logger.warning(f"File not found for upload: {filepath}")
            return

        with open(filepath, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            run_data = list(reader)
        
        if not run_data: 
            logger.info(f"File {filepath} is empty. Skipping upload.")
            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()
        logger.info(f"Successfully saved {os.path.basename(filepath)} to Supabase.")
    except Exception as e:
        logger.error(f"Failed to save output file to Supabase: {e}")

def run_web_analyst_mode(input_file_path, output_dir_override=None, max_events=None, output_filename=None, enable_fallback=False):
    print("=" * 60)
    print(f"MODE: Web Research Analyst (Crawl4AI) | Version: {APP_VERSION}")
    print("=" * 60)
    
    effective_output_dir = output_dir_override or OUTPUT_DIR
    
    run_timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    
    try:
        with open(input_file_path, 'r', encoding='utf-8') as f:
            races = json.load(f)
        races.sort(key=lambda x: x.get('Priority', 99))
        print(f"[SUCCESS] Found {len(races)} events to process from '{input_file_path}'.")
    except (FileNotFoundError, json.JSONDecodeError) as e:
        print(f"[ERROR] CONFIGURATION ERROR: Could not read '{input_file_path}'. Error: {e}")
        return

    if max_events is not None and max_events > 0:
        races = races[:max_events]
        print(f"[INFO] Limiting run to a maximum of {max_events} events.")

    grouped_races, failed_missions = {},[]
    for race in races:
        raw_type = race.get("Type")
        race_type = str(raw_type if raw_type else "Unknown").lower()
        
        if race_type not in grouped_races: grouped_races[race_type] = []
        grouped_races[race_type].append(race)

    agent = None
    type_to_filepath = {} 

    try:
        for race_type, race_list in grouped_races.items():
            schema_map = {
                "triathlon": TRIATHLON_SCHEMA, "run": RUNNING_SCHEMA, "running": RUNNING_SCHEMA, 
                "trail running": RUNNING_SCHEMA, "swimming": SWIMMING_SCHEMA, "swimathon": SWIMMING_SCHEMA, 
                "duathlon": DUATHLON_SCHEMA, "aquathlon": AQUATHLON_SCHEMA, "aquabike": AQUABIKE_SCHEMA, 
                "cycling": CYCLING_SCHEMA, "fitness racing": FITNESS_RACING_SCHEMA
            }
            schema = schema_map.get(race_type)
            if not schema:
                print(f"[WARNING] Skipping unknown race type '{race_type}'.")
                continue

            # Pass the fallback flag to the agent
            agent = MistralAnalystAgent(
                mistral_key=MISTRAL_API_KEY, 
                search_key=SEARCH_API_KEY, cse_id=CSE_ID, schema=schema,
                enable_fallback=enable_fallback
            )
            
            base_filename = f"{output_filename}_{race_type}" if output_filename else f"Crawl4AI_WebAnalyst_{race_type}_{run_timestamp}"
            final_filename = f"{base_filename}.csv"
            output_filepath = os.path.join(effective_output_dir, final_filename)
            type_to_filepath[race_type] = output_filepath

            print(f"\n[INFO] Processing {len(race_list)} '{race_type}' events. Target File: {final_filename}")
            
            with SafeCSVWriter(output_filepath, schema) as writer:
                for i, race_info in enumerate(race_list):
                    event_name = race_info.get("Festival")
                    if not event_name:
                        continue

                    print("\n" + "=" * 60)
                    print(f"[STARTING MISSION] {i + 1}/{len(race_list)} for '{race_type.upper()}': {event_name}")
                    print("=" * 60)

                    caching_key = agent.get_caching_key(event_name)
                    cache_file_path = os.path.join(KNOWLEDGE_CACHE_DIR, f"{caching_key}.json")
                    
                    knowledge_base = agent.run(race_info)

                    if knowledge_base:
                        with open(cache_file_path, 'w', encoding='utf-8') as f_cache:
                            json.dump(serialize_knowledge_base(knowledge_base), f_cache, indent=4)

                        for variant_name, data in knowledge_base.items():
                            if row := format_final_row(event_name, variant_name, data, schema):
                                writer.writerow(row)
                        print(f"[SUCCESS] Saved data for: {event_name}")
                    else:
                        print(f"[ERROR] No data extracted for: {event_name}")
                        failed_missions.append(event_name)

    except AgentInitializationError as e:
        print(f"[FATAL] Agent setup failed: {e}")
        # Not re-raising here to let the script exit gracefully instead of crashing streamlit
        return
    finally:
        if agent: agent.shutdown()
        
        print("\n[INFO] Syncing output files to database...")
        for r_type, filepath in type_to_filepath.items():
            if os.path.exists(filepath):
                save_output_to_supabase(filepath, agent_mode="web_analyst", event_type=r_type)

    print("\n" + "=" * 60)
    print("Web Research Analyst Run Complete")
    if failed_missions:
        print("\n[SUMMARY] Failed Missions:")
        for event in failed_missions: print(f"  - {event}")
    else:
        print("\n[SUCCESS] All missions completed successfully.")
    print("=" * 60)


def run_calendar_crawl_mode(output_dir_override=None, max_events=None, output_filename=None, enable_fallback=False):
    print("=" * 60)
    print(f"MODE: Calendar Crawl | Version: {APP_VERSION}")
    print("=" * 60)
    
    effective_output_dir = output_dir_override or OUTPUT_DIR
    schema = CYCLING_SCHEMA
    
    run_timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    agent = None
    failed_missions =[]

    try:
        agent = MistralAnalystAgent(
            mistral_key=MISTRAL_API_KEY,
            search_key=SEARCH_API_KEY, cse_id=CSE_ID, schema=schema,
            enable_fallback=enable_fallback
        )

        CALENDAR_URL = "https://www.audaxindia.in/events.php"
        main_page_content = agent._get_content_from_url(CALENDAR_URL, is_web_research=False)
        
        if not main_page_content:
            print(f"[FATAL] Could not retrieve content from the calendar URL. Aborting.")
            return

        event_list =[]
        organizer_pattern = re.compile(r"^\s*\*\*(?P<city>[^*]+)\*\*\s*(?P<name>[^\[\n\r]+)")
        link_pattern = re.compile(r'\[(?P<details>.*?)\]\((?P<url>.*?)\)')
        
        for line in main_page_content.splitlines():
            organizer_match = organizer_pattern.match(line)
            if organizer_match:
                city, name = organizer_match.group('city').strip(), organizer_match.group('name').strip()
                current_organizer = f"{city} {name}"
                for link_match in link_pattern.finditer(line):
                    event_list.append({
                        "organizer": current_organizer, 
                        "event_details": link_match.group('details').strip(),
                        "url": urljoin(CALENDAR_URL, link_match.group('url'))
                    })

        if not event_list:
            print("[INFO] No events found on calendar.")
            return
        
        print(f"[SUCCESS] Found {len(event_list)} total events.")
        
        if max_events is not None and max_events > 0:
            events_to_process = event_list[:max_events]
            print(f"[INFO] Limiting run to {max_events} events.")
        else:
            events_to_process = event_list

        base_filename = output_filename or f"CalendarCrawl_Cycling_{run_timestamp}"
        final_filename = f"{base_filename}.csv"
        output_filepath = os.path.join(effective_output_dir, final_filename)
        
        print(f"[INFO] Writing all output to: {output_filepath}")

        with SafeCSVWriter(output_filepath, schema) as writer:
            for i, event_info in enumerate(events_to_process):
                event_name = f"{event_info['organizer']} {event_info['event_details']}"
                event_url = event_info['url']
                print("\n" + "=" * 60)
                print(f"[STARTING MISSION] {i + 1}/{len(events_to_process)} FOR: {event_name}")
                print("=" * 60)

                race_info = {"Festival": event_name, "Type": "Cycling"}
                url_part = hashlib.md5(event_url.encode()).hexdigest()[:8]
                caching_key = agent.get_caching_key(f"{event_name}-{url_part}")
                cache_file_path = os.path.join(KNOWLEDGE_CACHE_DIR, f"{caching_key}.json")

                knowledge_base = agent.run_direct(race_info, direct_urls=[event_url])
                
                if knowledge_base:
                    with open(cache_file_path, 'w', encoding='utf-8') as f:
                        json.dump(serialize_knowledge_base(knowledge_base), f, indent=4)

                    main_festival_name = event_info['organizer']
                    for variant_name, data in knowledge_base.items():
                        if row := format_final_row(main_festival_name, variant_name, data, schema):
                            writer.writerow(row)
                    print(f"[SUCCESS] Saved data for: {event_name}")
                else:
                    print(f"[ERROR] No data found for: {event_name}")
                    failed_missions.append(event_name)
        
        save_output_to_supabase(output_filepath, agent_mode="calendar_crawl", event_type="Cycling")

    except AgentInitializationError as e:
        print(f"[FATAL] Error: {e}")
        return
    finally:
        if agent: agent.shutdown()

    print("\n" + "=" * 60)
    print("Calendar Crawl Complete")
    if failed_missions:
        print("\n[SUMMARY] Failed Missions:")
        for event in failed_missions: print(f"  - {event}")
    else:
        print("\n[SUCCESS] All missions completed successfully.")
    print("=" * 60)