# agent.py

# --- CRITICAL FIX FOR CHROMADB SQLITE VERSION ---
import sys
try:
    __import__('pysqlite3')
    sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
except ImportError:
    pass
import os
os.environ["ANONYMIZED_TELEMETRY"] = "False"
os.environ["CHROMA_TELEMETRY"] = "false"
# -----------------------------------------------

import re
import json
import time
import hashlib
from datetime import datetime, timezone
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import wraps

import requests
import dirtyjson
import chromadb
from chromadb.config import Settings
from sentence_transformers import CrossEncoder
from langchain_mistralai.chat_models import ChatMistralAI
from langchain_core.messages import HumanMessage
from markdown_it import MarkdownIt
from rank_bm25 import BM25Okapi
import numpy as np
import markdown
import trafilatura
from bs4 import BeautifulSoup
import spacy

from config import (
    MISTRAL_MODEL, MAX_RETRIES, DEBUG, MAX_CONCURRENT_CRAWLERS,
    TOP_N_URLS_TO_PROCESS, CRAWL_CACHE_DIR, SPACY_MODEL,
    VECTOR_DB_PATH, CROSS_ENCODER_MODEL, RAG_CANDIDATE_POOL_SIZE,
    RAG_FINAL_EVIDENCE_COUNT, SUPABASE_CONFIGURED,
    OPENROUTER_API_KEY, PERPLEXITY_SONAR_MODEL, QWEN_EMBEDDING_MODEL, LLAMA_VERIFIER_MODEL,
    MISTRAL_API_KEY, SEARCH_API_KEY, CSE_ID, SONAR_ANALYTICS_DIR
)
from schemas import (
    DEFAULT_BLANK_FIELDS, CHOICE_OPTIONS, INFERABLE_FIELDS, BLACKLISTED_DOMAINS
)
# Imported the new DB2 fetcher
from supabase_client import get_knowledge_cache, set_knowledge_cache, download_from_storage, upload_to_storage, fetch_races_db_fields

class AgentInitializationError(Exception): pass

def clean_markdown_with_trafilatura(md_text: str) -> str:
    try:
        md_text = re.sub(r'!\[.*?\]\(.*?\)', ' ', md_text)
        html = markdown.markdown(md_text)
        soup = BeautifulSoup(html, "html.parser")
        for tag in soup(["img", "nav", "footer", "header", "aside", "script", "style"]):
            tag.decompose()
        clean_html = str(soup)
        extracted = trafilatura.extract(clean_html, include_links=False, include_images=False, include_comments=False, favor_precision=True)
        if not extracted: return md_text 
        extracted = re.sub(r'\n{3,}', '\n\n', extracted)
        extracted = re.sub(r'[ \t]+', ' ', extracted)
        return extracted
    except Exception as e:
        print(f"  - [WARNING] Trafilatura text cleaning failed: {e}")
        return md_text

def retry(retries=MAX_RETRIES, delay=5):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            for i in range(retries):
                try:
                    return f(*args, **kwargs)
                except Exception as e:
                    if i < retries - 1:
                        time.sleep(delay * (2 ** i))
                    else:
                        print(f"  [ERROR] Function '{f.__name__}' failed after {retries} retries.");
                        return None
        return wrapper
    return decorator

class Field:
    def __init__(self, value=None, confidence=0.0, sources=None, inferred_by=""):
        self.value = value
        self.confidence = confidence
        self.sources = sources or[]
        self.inferred_by = inferred_by
        self.last_updated = datetime.now(timezone.utc).isoformat()

    def to_dict(self):
        return {
            "value": self.value, "confidence": self.confidence,
            "sources": self.sources, "inferred_by": self.inferred_by, "last_updated": self.last_updated
        }

    @classmethod
    def from_dict(cls, data):
        return cls(value=data.get('value'), confidence=data.get('confidence', 0.0), sources=data.get('sources',[]), inferred_by=data.get('inferred_by', ''))

class OpenRouterEmbedder:
    def __init__(self, api_key, model_name):
        self.api_key = api_key
        self.model_name = model_name
        self.url = "https://openrouter.ai/api/v1/embeddings"
        
    def encode(self, texts):
        if not self.api_key:
            raise Exception("OPENROUTER_API_KEY is missing for embeddings.")
        if isinstance(texts, str): texts = [texts]
            
        embeddings =[]
        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
        
        for i in range(0, len(texts), 50):
            batch = texts[i:i+50]
            data = {"model": self.model_name, "input": batch}
            resp = requests.post(self.url, headers=headers, json=data, timeout=60)
            if resp.status_code == 200:
                data_list = resp.json().get("data",[])
                data_list = sorted(data_list, key=lambda x: x["index"])
                embeddings.extend([item["embedding"] for item in data_list])
            else:
                raise Exception(f"Embedding API Error: {resp.status_code} - {resp.text}")
        return np.array(embeddings)


class MistralAnalystAgent:
    def __init__(self, mistral_key: str, search_key: str, cse_id: str, schema: list, enable_fallback: bool = False):
        if not all([mistral_key, search_key, cse_id]):
            raise AgentInitializationError("One or more API keys (Mistral, Search) are missing.")

        self.enable_fallback = enable_fallback
        self.llm_client = ChatMistralAI(api_key=mistral_key, model=MISTRAL_MODEL, temperature=0.0)
        self.search_api_key = search_key
        self.cse_id = cse_id
        self.schema = schema
        
        # Will raise error if missing context and fallback disabled
        self.field_instructions = self._generate_field_instructions()
        
        self.invalid_years =[str(y) for y in range(2015, 2025)]

        print("[INFO] Initializing ML models and VectorDB...")
        try:
            self.nlp = spacy.load(SPACY_MODEL)
        except (OSError, IOError) as e:
            raise AgentInitializationError(f"[FATAL] Failed to load spaCy model '{SPACY_MODEL}'. Error: {e}")

        self.chroma_client = chromadb.Client(Settings(persist_directory=VECTOR_DB_PATH, anonymized_telemetry=False, is_persistent=True, allow_reset=True))
        self.embedding_model = OpenRouterEmbedder(OPENROUTER_API_KEY, QWEN_EMBEDDING_MODEL)
        self.cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL)
        self.md_parser = MarkdownIt()
        self.chroma_collection = None
        self.bm25_index = None
        self.mission_corpus =[] 
        self.corpus_map = {} 
        self.mission_inference_cache = {}
        print("  -[SUCCESS] Models initialized.")

    def shutdown(self):
        try: self.chroma_client.reset()
        except: pass

    def get_legacy_caching_key(self, event_name: str) -> str:
        base_name = re.sub(r'sprint|standard|olympic|full iron|half iron|70\.3', '', event_name, flags=re.IGNORECASE)
        return re.sub(r'[^a-z0-9]+', '-', base_name.lower()).strip('-')

    def get_caching_key(self, event_name: str, url: str = "") -> str:
        if not event_name: event_name = "unknown_event"
        base_name = re.sub(r'sprint|standard|olympic|full iron|half iron|70\.3', '', event_name, flags=re.IGNORECASE)
        slug = re.sub(r'[^a-z0-9]+', '-', base_name.lower()).strip('-')
        name_hash = hashlib.md5(slug.encode('utf-8')).hexdigest()[:8]
        url_hash = ""
        if url: url_hash = "-" + hashlib.md5(url.encode('utf-8')).hexdigest()[:8]
        safe_slug = slug[:40].strip('-')
        if not safe_slug: safe_slug = "event"
        return f"{safe_slug}-{name_hash}{url_hash}"

    def _generate_field_instructions(self) -> dict:
        instructions = {}
        db_fields = fetch_races_db_fields()
        
        # ENFORCING FALLBACK LOGIC
        if not db_fields and not self.enable_fallback:
            raise AgentInitializationError("Failed to fetch schema fields from Database2 and Fallback is Disabled. Check Database2 settings.")

        meta_map = {row['field']: row for row in db_fields} if db_fields else {}

        for key in self.schema:
            if key in DEFAULT_BLANK_FIELDS: continue
            
            meta = meta_map.get(key)
            if meta:
                inst = f"Extract '{key}'"
                if meta.get('display_name'): inst += f" (also referred to as '{meta['display_name']}')."
                else: inst += "."
                
                if meta.get('question_text'): inst += f" Context constraint: {meta['question_text']}."
                
                opts = meta.get('data_options')
                if opts:
                    if isinstance(opts, str):
                        if opts.strip().startswith('[') and opts.strip().endswith(']'):
                            try:
                                opt_list = json.loads(opts)
                                opt_str = ', '.join([str(o) for o in opt_list])
                            except: opt_str = opts.replace('\n', ', ')
                        else: opt_str = opts.replace('\n', ', ')
                    else: opt_str = str(opts)
                    inst += f" MUST be exactly one of: [{opt_str}]."
                elif key in CHOICE_OPTIONS: 
                    inst += f" MUST be exactly one of: {', '.join(CHOICE_OPTIONS[key])}."
                
                dt = meta.get('data_type')
                df = meta.get('data_format')
                if dt or df: inst += f" Required format/type: {dt or ''} {df or ''}."
                
                instructions[key] = inst.strip()
            else:
                # ENFORCING FALLBACK LOGIC
                if not self.enable_fallback:
                    raise AgentInitializationError(f"Field '{key}' not found in Database2 schema and Fallback is Disabled.")

                if key in CHOICE_OPTIONS: instructions[key] = f"Extract '{key}'. MUST be one of: {', '.join(CHOICE_OPTIONS[key])}."
                else: instructions[key] = f"Extract '{key}'."
                    
        return instructions

    def _normalize_key(self, key_str: str) -> str:
        return re.sub(r'[^a-z0-9]', '', str(key_str).lower())

    def _clean_citations(self, text: str) -> str:
        if not isinstance(text, str): return text
        return re.sub(r'\[\d+\]', '', text).strip()

    def _clean_dirtyjson(self, obj):
        if isinstance(obj, dict):
            return {self._clean_citations(k): self._clean_dirtyjson(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return[self._clean_dirtyjson(v) for v in obj]
        elif isinstance(obj, str):
            return self._clean_citations(obj)
        else:
            return obj

    def _parse_json_response(self, response_text: str):
        if not response_text: return None
        clean_text = self._clean_citations(response_text)
        match = re.search(r'\{.*\}|\[.*\]', clean_text, re.DOTALL)
        if not match: return None
        raw = match.group(0)
        try:
            return self._clean_dirtyjson(json.loads(raw))
        except json.JSONDecodeError:
            try:
                parsed = dirtyjson.loads(raw)
                return self._clean_dirtyjson(parsed)
            except Exception:
                return None

    @retry()
    def _call_llm(self, prompt: str) -> str:
        return self.llm_client.invoke([HumanMessage(content=prompt)]).content

    @retry()
    def _call_openrouter_gap_filler(self, prompt: str) -> str:
        if not OPENROUTER_API_KEY: return None
        headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
        data = {
            "model": PERPLEXITY_SONAR_MODEL,
            "messages":[{"role": "user", "content": prompt}],
            "temperature": 0.1,
            "max_tokens": 2000
        }
        try:
            resp = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=data, timeout=60)
            if resp.status_code == 200: return resp.json()['choices'][0]['message']['content']
            else:
                print(f"  - [ERROR] Perplexity API Error: {resp.status_code} - {resp.text}")
                return None
        except Exception as e:
            print(f"  -[ERROR] Perplexity connection failed: {e}")
            return None

    @retry()
    def _call_openrouter_llama(self, prompt: str) -> str:
        if not OPENROUTER_API_KEY: return None
        headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
        data = {
            "model": LLAMA_VERIFIER_MODEL,
            "messages":[{"role": "user", "content": prompt}],
            "temperature": 0.0,
            "max_tokens": 3000
        }
        try:
            resp = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=data, timeout=60)
            if resp.status_code == 200: return resp.json()['choices'][0]['message']['content']
            else:
                print(f"  -[ERROR] Llama API Error: {resp.status_code} - {resp.text}")
                return None
        except Exception as e:
            print(f"  - [ERROR] Llama connection failed: {e}")
            return None

    def _step_1a_initial_search(self, race_info: dict) -> list:
        event_name = race_info.get("Festival")
        print(f"\n[STEP 1A] Performing initial search for '{event_name}'")
        query = f'{event_name} 2025 OR 2026'
        try:
            url = "https://www.googleapis.com/customsearch/v1"
            params = {"key": self.search_api_key, "cx": self.cse_id, "q": query, "num": 10}
            response = requests.get(url, params=params)
            response.raise_for_status()
            
            raw_results =[{"title": i.get("title"), "link": i.get("link"), "snippet": i.get("snippet")} for i in response.json().get("items", [])]
            clean_results =[r for r in raw_results if not any(d in r['link'].lower() for d in BLACKLISTED_DOMAINS)]
            return clean_results
        except requests.HTTPError as e: return[]

    def _step_1b_validate_and_select_urls(self, event_name: str, search_results: list) -> list:
        print("[STEP 1B] Validating search results with LLM...")
        if not search_results: return[]

        prompt = f"Identify the most relevant websites for '{event_name}'. Select the single best 'primary_url' (official page) and up to three 'secondary_urls'.\n\nSearch Results:\n{json.dumps(search_results, indent=2)}\n\nResponse MUST be JSON: {{'primary_url': '...', 'secondary_urls': [...]}}"
        response_text = self._call_llm(prompt)
        parsed = self._parse_json_response(response_text)
        if parsed and isinstance(parsed, dict):
            primary = parsed.get("primary_url")
            secondaries = parsed.get("secondary_urls",[])
            final_urls = list(dict.fromkeys([u for u in ([primary] + secondaries if primary else secondaries) if u]))
            return final_urls[:TOP_N_URLS_TO_PROCESS]
        
        if response_text:
            salvaged_urls = re.findall(r'https?://[^\s"\'\)\],]+', response_text)
            if salvaged_urls: 
                return list(dict.fromkeys(salvaged_urls))[:TOP_N_URLS_TO_PROCESS]
                
        return [r['link'] for r in search_results[:TOP_N_URLS_TO_PROCESS] if r.get('link')]

    @retry(retries=2, delay=10)
    def _get_content_from_url(self, url: str, is_web_research: bool) -> str | None:
        if not url.strip(): return None
        url_hash = hashlib.md5((url + str(is_web_research)).encode()).hexdigest()
        storage_file_path = f"{url_hash}.md"
        
        if SUPABASE_CONFIGURED:
            content = download_from_storage(storage_file_path)
            if content: return content
            
        local_cache_path = os.path.join(CRAWL_CACHE_DIR, f"{url_hash}.md")
        if os.path.exists(local_cache_path):
            with open(local_cache_path, 'r', encoding='utf-8') as f: return f.read()
            
        print(f"  - No cache found. Crawling: {url}")
        try:
            api_url = f"https://r.jina.ai/{url}"
            response = requests.get(api_url, timeout=60)
            if response.status_code == 200 and response.text:
                content = response.text
                if is_web_research:
                    print("    - Applying Trafilatura semantic extraction...")
                    content = clean_markdown_with_trafilatura(content)
                with open(local_cache_path, 'w', encoding='utf-8') as f: f.write(content)
                if SUPABASE_CONFIGURED: upload_to_storage(storage_file_path, content)
                return content
            return None
        except requests.RequestException: return None

    def _chunk_and_index_text(self, text: str, url: str, event_id_str: str):
        if not text: return
        from langchain_text_splitters import RecursiveCharacterTextSplitter
        chunks = RecursiveCharacterTextSplitter(chunk_size=768, chunk_overlap=100).split_text(text)
        if not chunks: return
        
        chunk_ids =[f"{event_id_str}_{hashlib.md5(chunk.encode()).hexdigest()}" for chunk in chunks]
        unique_chunk_ids = list(set(chunk_ids))
        existing_chunks = self.chroma_collection.get(ids=unique_chunk_ids)
        existing_ids = set(existing_chunks['ids'])
        
        new_chunks_to_add =[]
        start_idx = len(self.mission_corpus)

        for i, chunk_id in enumerate(chunk_ids):
            self.mission_corpus.append(chunks[i])
            self.corpus_map[start_idx + i] = {'id': chunk_id, 'snippet': chunks[i]}
            if chunk_id not in existing_ids:
                new_chunks_to_add.append({'id': chunk_id, 'chunk': chunks[i]})
                existing_ids.add(chunk_id)
        
        if new_chunks_to_add:
            print(f"    - Indexing {len(new_chunks_to_add)} new passages from {url}")
            new_ids = [item['id'] for item in new_chunks_to_add]
            new_documents =[item['chunk'] for item in new_chunks_to_add]
            try:
                new_embeddings = self.embedding_model.encode(new_documents).tolist()
                new_metadatas =[{"source_url": url, "event_id": event_id_str} for _ in new_ids]
                if new_ids:
                    self.chroma_collection.add(ids=new_ids, embeddings=new_embeddings, documents=new_documents, metadatas=new_metadatas)
            except Exception as e:
                print(f"    -[ERROR] Failed to embed documents: {e}")

    def _build_bm25_index(self):
        if self.mission_corpus:
            tokenized_corpus =[doc.lower().split() for doc in self.mission_corpus]
            self.bm25_index = BM25Okapi(tokenized_corpus)

    def _retrieve_and_fuse_evidence(self, query: str, top_k: int) -> list[dict]:
        collection_count = self.chroma_collection.count()
        if collection_count == 0: return[]
        
        vector_hits = {}
        try:
            query_embedding = self.embedding_model.encode([query]).tolist()
            n_results = min(top_k * 2, collection_count)
            chroma_results = self.chroma_collection.query(query_embeddings=query_embedding, n_results=n_results)
            if chroma_results['ids']:
                for rank, (doc_id, doc) in enumerate(zip(chroma_results['ids'][0], chroma_results['documents'][0])):
                    vector_hits[doc_id] = {'snippet': doc, 'rank': rank}
        except Exception: pass

        keyword_hits = {}
        if self.bm25_index:
            tokenized_query = query.lower().split()
            doc_scores = self.bm25_index.get_scores(tokenized_query)
            top_indices = sorted(range(len(doc_scores)), key=lambda i: doc_scores[i], reverse=True)[:top_k*2]
            for rank, idx in enumerate(top_indices):
                if idx in self.corpus_map:
                    item = self.corpus_map[idx]
                    keyword_hits[item['id']] = {'snippet': item['snippet'], 'rank': rank}

        fused_scores = {}
        k = 60
        for doc_id, hit in vector_hits.items():
            if doc_id not in fused_scores: fused_scores[doc_id] = 0
            fused_scores[doc_id] += 1 / (k + hit['rank'] + 1)
        for doc_id, hit in keyword_hits.items():
            if doc_id not in fused_scores: fused_scores[doc_id] = 0
            fused_scores[doc_id] += 1 / (k + hit['rank'] + 1)
            
        sorted_ids = sorted(fused_scores.keys(), key=lambda x: fused_scores[x], reverse=True)[:top_k]
        final_results =[]
        for doc_id in sorted_ids:
            if doc_id in vector_hits: final_results.append({'id': doc_id, 'snippet': vector_hits[doc_id]['snippet']})
            elif doc_id in keyword_hits: final_results.append({'id': doc_id, 'snippet': keyword_hits[doc_id]['snippet']})
        return final_results

    def _rerank_evidence_with_cross_encoder(self, query: str, evidence: list[dict]) -> list[dict]:
        if not evidence: return []
        pairs =[(query, item['snippet']) for item in evidence]
        scores = self.cross_encoder.predict(pairs)
        for i, item in enumerate(evidence): 
            item['rerank_score'] = float(scores[i])
        return sorted(evidence, key=lambda x: x['rerank_score'], reverse=True)

    def _perform_holistic_extraction(self, knowledge_base: dict, event_name: str, variant_name: str):
        print(f"    - Pass 1: Holistic RAG extraction for '{variant_name}'...")
        all_fields_query = f"Extract all available information for the event '{event_name}', specifically for the '{variant_name}' category."
        candidate_evidence = self._retrieve_and_fuse_evidence(all_fields_query, top_k=RAG_CANDIDATE_POOL_SIZE)
        reranked_evidence = self._rerank_evidence_with_cross_encoder(all_fields_query, candidate_evidence)
        final_evidence = reranked_evidence[:RAG_FINAL_EVIDENCE_COUNT + 5]

        if not final_evidence: return knowledge_base

        evidence_prompt = "\n".join([f"Evidence Snippet:\n---\n{e['snippet']}\n---" for e in final_evidence])
        fields_to_extract =[f"- {key}: {desc}" for key, desc in self.field_instructions.items()]
        fields_prompt = "\n".join(fields_to_extract)

        prompt = f"""You are a strict data extraction assistant. Based on the evidence below, extract all specified fields for '{variant_name}' of '{event_name}'.
CRITICAL RULES:
- Extract ONLY verifiable information directly from the provided source.
- DO NOT infer, assume, estimate, or create missing values.
- If a value is not explicitly present, omit the key or set it to null.
- **IMPORTANT**: For prices, distances, and cutoffs, ensure they match the EXACT variant '{variant_name}'. If the evidence shows data for a DIFFERENT distance, return null.
- Return ONLY flat strings. NO nested JSON/dictionaries inside values.

## Evidence
{evidence_prompt}

## Task
Extract these fields into a JSON object:
{fields_prompt}

## Response
JSON object only."""

        response_text = self._call_llm(prompt)
        extracted_data = self._parse_json_response(response_text)
        
        if extracted_data and isinstance(extracted_data, dict):
            norm_schema_keys = {self._normalize_key(k): k for k in self.schema}
            filled = 0
            confidence = 0.85
            for key, value in extracted_data.items():
                norm_key = norm_schema_keys.get(self._normalize_key(key))
                if norm_key and value and str(value).lower() not in["null", "none", "unknown", ""]:
                    if not knowledge_base[variant_name].get(norm_key, Field()).value:
                        val_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
                        knowledge_base[variant_name][norm_key] = Field(value=val_str, confidence=confidence, inferred_by="rag_holistic", sources=final_evidence)
                        filled += 1
            print(f"      - [SUCCESS] Holistic extracted {filled} fields.")
        return knowledge_base

    def _perform_targeted_recovery(self, knowledge_base: dict, event_name: str, variant_name: str):
        print(f"    - Pass 2: Targeted Recovery for missing fields...")
        missing_fields =[k for k in self.schema if k not in DEFAULT_BLANK_FIELDS and not knowledge_base[variant_name].get(k, Field()).value]
        
        if not missing_fields: return knowledge_base
        
        for field in missing_fields:
            query = f"What is the {field} for {variant_name} in {event_name}?"
            evidence = self._retrieve_and_fuse_evidence(query, top_k=10)
            reranked = self._rerank_evidence_with_cross_encoder(query, evidence)[:3]
            
            if not reranked: continue
            evidence_text = "\n".join([e['snippet'] for e in reranked])
            
            instruction = self.field_instructions.get(field, f"Extract '{field}'.")

            prompt = f"""Based on the evidence, find the exact value for '{field}' for the race '{variant_name}'.
{instruction}

CRITICAL RULES:
- Extract ONLY verifiable information directly from the provided source.
- Ensure the value applies EXACTLY to the '{variant_name}'. If the evidence is for a different distance/variant, return null.
- Return flat strings only. No dicts.

Evidence: {evidence_text}
Return a JSON object: {{"answer": "value"}} or {{"answer": null}} if not found."""
            
            response = self._call_llm(prompt)
            parsed = self._parse_json_response(response)
            if parsed and isinstance(parsed, dict):
                ans = parsed.get('answer')
                if ans and str(ans).lower() not in['null', 'none', 'unknown', '']:
                    val_str = json.dumps(ans) if isinstance(ans, (dict, list)) else str(ans)
                    knowledge_base[variant_name][field] = Field(value=val_str, confidence=0.9, inferred_by="rag_targeted", sources=reranked)
        return knowledge_base

    def _perform_batched_gap_fill(self, knowledge_base: dict, event_name: str):
        if not OPENROUTER_API_KEY: return knowledge_base

        print(f"    - Pass 3: Batched Gap Filling ({PERPLEXITY_SONAR_MODEL})...")
        
        missing_data_map = {}
        for variant, fields in knowledge_base.items():
            missing_fields =[k for k in self.schema if k not in DEFAULT_BLANK_FIELDS and not fields.get(k, Field()).value]
            if missing_fields: missing_data_map[variant] = missing_fields

        if not missing_data_map: return knowledge_base

        records_str = ""
        all_missing_fields_set = set()
        
        for i, (variant, m_fields) in enumerate(missing_data_map.items()):
            field_list_str = "\n   - ".join(m_fields)
            records_str += f"\n{i+1}. Record ID: \"{variant}\"\n   Missing fields:\n   - {field_list_str}\n"
            for f in m_fields: all_missing_fields_set.add(f)

        fields_to_extract =[f"- {key}: {self.field_instructions.get(key, 'Extract exact value')}" for key in all_missing_fields_set]
        field_specific_instructions = "\n".join(fields_to_extract)

        prompt = f"""Search the live web to fill the following missing fields for the athletic event '{event_name}'.

Records:
{records_str}

Search Methodology:
- Use your ability to search the web to find official, accurate data for this specific event.
- Do not repeat or re-fetch fields not listed.

CRITICAL Extraction Rules:
- DIFFERENT VARIANTS HAVE DIFFERENT PRICES AND CUTOFFS. You MUST extract the specific registrationCost and Cutoff time for EACH individual variant. Do not blindly copy one variant's fee to another.
- Extract ONLY verifiable information found online.
- DO NOT infer, assume, estimate, or create missing values.
- If specific information is unavailable on the web, return: NULL
- Return ONLY flat strings. NO nested objects!

Field Definitions (Strictly follow these definitions):
{field_specific_instructions}

Output Format:
Return strict JSON matching exactly this schema:
{{
  "<record_id_exactly_as_written_above>": {{
    "<field_name_exactly_as_written>": "<exact_value_or_NULL>"
  }}
}}"""
        response_text = self._call_openrouter_gap_filler(prompt)
        
        if response_text:
            os.makedirs(SONAR_ANALYTICS_DIR, exist_ok=True)
            clean_name = re.sub(r'[^a-zA-Z0-9]+', '_', event_name[:20]).strip('_')
            analytics_file = os.path.join(SONAR_ANALYTICS_DIR, f"sonar_gap_{clean_name}_{int(time.time())}.json")
            with open(analytics_file, 'w', encoding='utf-8') as f:
                f.write(response_text)

            filled_data = self._parse_json_response(response_text)
            
            if filled_data and isinstance(filled_data, dict):
                filled_count = 0
                norm_kb_keys = {self._normalize_key(k): k for k in knowledge_base.keys()}
                norm_schema_keys = {self._normalize_key(k): k for k in self.schema}

                for variant, fields in filled_data.items():
                    norm_var = self._normalize_key(variant)
                    actual_var = norm_kb_keys.get(norm_var)
                    
                    if actual_var and isinstance(fields, dict):
                        for key, value in fields.items():
                            norm_key = norm_schema_keys.get(self._normalize_key(key))
                            
                            if norm_key and value and str(value).lower() not in["null", "none", "unknown", ""]:
                                val_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
                                knowledge_base[actual_var][norm_key] = Field(value=val_str, confidence=0.95, inferred_by="sonar_gap_fill")
                                filled_count += 1
                print(f"      - [SUCCESS] Perplexity Sonar filled {filled_count} missing gaps.")
        return knowledge_base

    def _verify_hallucinations(self, knowledge_base: dict):
        if not OPENROUTER_API_KEY: return knowledge_base
        print(f"    - Pass 4: Hallucination Check ({LLAMA_VERIFIER_MODEL})...")
        
        kb_json = {}
        for variant, fields in knowledge_base.items():
            kb_json[variant] = {}
            for k, f in fields.items():
                if f.value and k not in DEFAULT_BLANK_FIELDS: 
                    kb_json[variant][k] = f.value
        
        if not kb_json: return knowledge_base

        prompt = f"""You are a strict Data Integrity Verifier. 
Review the following extracted event data for hallucinations, logical impossibilities, or invalid formats.

Data:
{json.dumps(kb_json, indent=2)}

Rules for Rejection (Set to NULL if violated):
- Dates must be logically possible.
- Prices must be realistic.
- Distances must match reality (e.g., ultras can be 50k, 75k, 100k, 100M). Do not reject standard ultra distances.
- Text fields must not contain conversational filler (e.g., "The price is $50").

Return ONLY a JSON object with the exact same structure. Replace any hallucinated or invalid values with NULL. Do not add any extra text or markdown outside the JSON.
"""
        response = self._call_openrouter_llama(prompt)
        if response:
            verified_data = self._parse_json_response(response)
            
            if verified_data and isinstance(verified_data, dict):
                rejected_count = 0
                
                norm_kb_keys = {self._normalize_key(k): k for k in knowledge_base.keys()}
                norm_schema_keys = {self._normalize_key(k): k for k in self.schema}
                
                for variant, fields in verified_data.items():
                    actual_var = norm_kb_keys.get(self._normalize_key(variant))
                    if actual_var and isinstance(fields, dict):
                        for key, val in fields.items():
                            actual_key = norm_schema_keys.get(self._normalize_key(key))
                            if actual_key and (val is None or str(val).lower() == "null"):
                                knowledge_base[actual_var][actual_key] = Field() 
                                rejected_count += 1
                if rejected_count > 0:
                    print(f"      - [WARNING] Llama rejected {rejected_count} fields as hallucinations.")
                else:
                    print(f"      -[SUCCESS] Llama verified data integrity.")
        return knowledge_base

    def _discover_and_validate_variants_from_content(self, content: str, knowledge_base: dict):
        print("    - Discovering all race variants from content...")
        prompt = f"""You are a strict data extractor. Identify the specific, official race categories or distances explicitly available for registration in the text provided.
CRITICAL RULES:
1. Extract ONLY variants explicitly mentioned in the text.
2. DO NOT invent, assume, or infer distances (e.g., do not add "Half Marathon" if only "50K" is listed).
3. DO NOT list general distances, only those meant for this specific event.
Return ONLY a JSON list of strings. If no variants are found, return an empty list[].
Text: {content[:15000]}..."""
        
        response_text = self._call_llm(prompt)
        variants = self._parse_json_response(response_text)
        
        if isinstance(variants, list):
            cleaned_variants =[v for v in variants if isinstance(v, str) and len(v) < 100]
            if cleaned_variants:
                print(f"      - [SUCCESS] Found variants: {', '.join(cleaned_variants)}")
                for variant_name in cleaned_variants:
                    if variant_name not in knowledge_base:
                        knowledge_base[variant_name] = {field: Field() for field in self.schema}
                return
        print("      - [WARNING] Variant discovery failed. Using default.")

    def _inject_pre_filled_data(self, knowledge_base: dict, pre_filled_data: dict, pre_filled_confidence: float = 0.3):
        if not pre_filled_data: return
        print(f"    - Injecting pre-filled data (confidence={pre_filled_confidence})...")
        for variant_name in knowledge_base.keys():
            for key, value in pre_filled_data.items():
                if key in self.schema and value:
                    if knowledge_base[variant_name].get(key, Field()).confidence < pre_filled_confidence:
                        val_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
                        knowledge_base[variant_name][key] = Field(value=val_str, confidence=pre_filled_confidence, inferred_by="pre_processed_data")

    def _run_inferential_filling(self, knowledge_base: dict):
        print("\n[INFERENCE] Running final analysis to infer missing data...")
        for variant, data in knowledge_base.items():
            city = data.get('city', Field()).value
            country = data.get('country', Field()).value
            if city and not country:
                resp = self._call_llm(f"What country is {city} in? Return ONLY the country name.")
                if resp:
                    knowledge_base[variant]['country'] = Field(value=resp.strip(), confidence=0.8, inferred_by="inference")
        return knowledge_base

    def determine_event_type(self, url: str) -> str | None:
        print(f"  - Determining event type for URL: {url}")
        content = self._get_content_from_url(url, is_web_research=False)
        if not content: return None
        valid_types = ", ".join(CHOICE_OPTIONS["type"])
        prompt = f"Analyze the text and determine the athletic event type. Answer MUST be one of: {valid_types}.\nText: {content[:2000]}"
        response = self._call_llm(prompt)
        if not response: return None
        for et in CHOICE_OPTIONS["type"]:
            if et.lower() in response.lower():
                print(f"    - [SUCCESS] Determined type: {et}")
                return et
        return None

    def _check_supabase_for_cache(self, event_key: str, legacy_key: str = None) -> dict | None:
        if not SUPABASE_CONFIGURED: return None
        data = get_knowledge_cache(event_key)
        if data:
            print(f"[SUCCESS] Found cache for '{event_key}' in Supabase.")
            return data
        if legacy_key and legacy_key != event_key:
            print(f"  - Primary key missing. Checking legacy key: '{legacy_key}'...")
            data = get_knowledge_cache(legacy_key)
            if data:
                print(f"[SUCCESS] Found legacy cache for '{legacy_key}' in Supabase.")
                return data
        return None

    def run(self, race_info: dict) -> dict | None:
        """Called by Web Research Agent."""
        event_name = race_info.get("Festival")
        event_key = self.get_caching_key(event_name)
        legacy_key = self.get_legacy_caching_key(event_name)
        
        cached_data = self._check_supabase_for_cache(event_key, legacy_key)
        if cached_data:
            return {v: {f: Field.from_dict(d) for f, d in fs.items()} for v, fs in cached_data.items()}
        
        search_results = self._step_1a_initial_search(race_info)
        if not search_results: return None
        validated_urls = self._step_1b_validate_and_select_urls(event_name, search_results)
        if not validated_urls: return None
        
        knowledge_base = self._crawl_and_extract(validated_urls, race_info, is_web_research=True)
        
        if knowledge_base and SUPABASE_CONFIGURED:
            serializable_kb = {v: {f: field.to_dict() for f, field in fs.items()} for v, fs in knowledge_base.items()}
            set_knowledge_cache(event_key, serializable_kb)
        return knowledge_base

    def run_direct(self, race_info: dict, direct_urls: list, pre_filled_data: dict = None, pre_filled_confidence: float = 0.3) -> dict | None:
        """Called by Prescraped Data Processor."""
        event_name = race_info.get("Festival")
        url_part = hashlib.md5(direct_urls[0].encode()).hexdigest()[:8]
        event_key = self.get_caching_key(f"{event_name}-{url_part}")
        legacy_key = self.get_legacy_caching_key(event_name)
        
        cached_data = self._check_supabase_for_cache(event_key, legacy_key)
        if cached_data:
            return {v: {f: Field.from_dict(d) for f, d in fs.items()} for v, fs in cached_data.items()}

        validated_urls =[url for url in direct_urls if self._is_valid_url(url)]
        if not validated_urls: return None
        
        knowledge_base = self._crawl_and_extract(validated_urls, race_info, pre_filled_data, pre_filled_confidence, is_web_research=False)
        
        if knowledge_base and SUPABASE_CONFIGURED:
            serializable_kb = {v: {f: field.to_dict() for f, field in fs.items()} for v, fs in knowledge_base.items()}
            set_knowledge_cache(event_key, serializable_kb)
        return knowledge_base

    def _is_valid_url(self, url: str) -> bool:
        if not url or not url.startswith(('http://', 'https://')): return False
        if any(year in url for year in self.invalid_years): return False
        return True

    def _crawl_and_extract(self, urls: list, race_info: dict, pre_filled_data: dict = None, pre_filled_confidence: float = 0.3, is_web_research: bool = False) -> dict:
        event_name = race_info.get("Festival")
        event_id_str = self.get_caching_key(event_name)
        self.chroma_collection = self.chroma_client.get_or_create_collection(name=f"coll_{event_id_str}")
        print(f"\n[STEP 2] Starting RAG processing for '{event_name}' (Collection: coll_{event_id_str})")
        
        self.mission_corpus =[]
        self.corpus_map = {}
        self.bm25_index = None

        knowledge_base = {}
        all_content =[]
        with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_CRAWLERS) as executor:
            futures = {executor.submit(self._get_content_from_url, url, is_web_research): url for url in urls}
            for future in as_completed(futures):
                if content := future.result():
                    all_content.append(content)
                    self._chunk_and_index_text(content, futures[future], event_id_str)

        if not all_content:
            print("  -[WARNING] No content crawled.")
            return {}

        self._build_bm25_index()

        combined_content = "\n\n".join(all_content)
        self._discover_and_validate_variants_from_content(combined_content, knowledge_base)

        if not knowledge_base:
            knowledge_base[event_name] = {field: Field() for field in self.schema}

        if pre_filled_data:
            self._inject_pre_filled_data(knowledge_base, pre_filled_data, pre_filled_confidence)

        for variant in list(knowledge_base.keys()):
            self._perform_holistic_extraction(knowledge_base, event_name, variant)
            self._perform_targeted_recovery(knowledge_base, event_name, variant)

        self._perform_batched_gap_fill(knowledge_base, event_name)
        self._verify_hallucinations(knowledge_base)
        self._run_inferential_filling(knowledge_base)
        
        print("\n[SUCCESS] Extraction complete.")
        return knowledge_base