"""
Automated audit script for Inframat-X RAG chatbot.
Evaluates Hit Rate@8 (At least one correct document found).
"""
import os
import re
import json
import time
import pandas as pd
from datetime import datetime
from typing import Tuple, Optional, Callable

def load_sources_map(csv_path="sources.csv"):
    if not os.path.exists(csv_path):
        return {}
    df = pd.read_csv(csv_path).fillna("")
    df.columns = df.columns.str.strip()
    src_map = {}
    for _, r in df.iterrows():
        raw_key = str(r.get("source_key", "")).strip().lower()
        fname = os.path.basename(raw_key).lower().strip()
        raw_name = str(r.get("name", "")).strip().lower()
        raw_id = str(r.get("id", "")).strip()
        
        clean_id = raw_id.replace("PAPER_", "").replace("paper_", "").lstrip("0")
        if not clean_id: clean_id = "0"
        
        if fname: src_map[fname.replace('.pdf', '')] = clean_id
        if raw_name: src_map[raw_name.replace('.pdf', '')] = clean_id
        src_map[raw_id.lower()] = clean_id
    return src_map

def extract_retrieved_ids(full_output: str) -> list:
    if not full_output:
        return []
    sources_match = re.search(r'\*\*Sources:\*\*(.*)', full_output)
    if sources_match:
        ids = re.findall(r'\[(\d+)\]', sources_match.group(1))
        return list(set(ids))
    ref_section = re.search(r'### References\s*\n(.*?)(?:\n\s*\n|$)', full_output, re.DOTALL)
    if ref_section:
        ids = re.findall(r'\[(\d+)\]', ref_section.group(1))
        return list(set(ids))
    return []

def calculate_hit_rate(retrieved_ids: list, gold_docs: list, sources_map: dict) -> float:
    """
    Checks if AT LEAST ONE expected document was successfully retrieved.
    Returns 1.0 (Success) or 0.0 (Fail).
    """
    if not gold_docs:
        return 0.0
        
    expected_ids = set()
    for g in gold_docs:
        g_clean = g.lower().replace('.pdf', '').strip()
        if g_clean in sources_map:
            expected_ids.add(sources_map[g_clean])
        else:
            nums = re.findall(r'\d+', g_clean)
            if nums:
                expected_ids.add(nums[-1].lstrip('0') or '0')
                
    # YOUR LOGIC: Did we find at least one?
    for e in expected_ids:
        if e in retrieved_ids:
            return 1.0  # 100% Success for this question
            
    return 0.0  # 0% Success

def run_audit(
    rag_reply_func, 
    gold_csv_path: str = "gold.csv",
    output_base_dir: Optional[str] = None,
    progress_callback: Optional[Callable[[str, int, int], None]] = None,
    k_retrieval: int = 10
) -> Tuple[str, str]:
    
    if not os.path.exists(gold_csv_path):
        return f"❌ Error: Could not find {gold_csv_path}.", ""
            
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    if output_base_dir is None:
        output_base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), f"Audit_{timestamp}")
    os.makedirs(output_base_dir, exist_ok=True)
        
    df = pd.read_csv(gold_csv_path)
    total_questions = len(df)
    jsonl_path = os.path.join(output_base_dir, "rag_logs.jsonl")
    
    sources_map = load_sources_map("sources.csv")
        
    total_hit_rate = 0.0
    processed_count = 0
        
    if progress_callback: progress_callback("Gold Set Benchmark", 0, total_questions)

    with open(jsonl_path, "w", encoding="utf-8") as log_file:
        for idx, row in df.iterrows():
            question = row['question']
            raw_gold = str(row['relevant_docs']).split(';')
            gold_docs = [p.strip() for p in raw_gold if p.strip()]
                        
            raw_output = rag_reply_func(question, k=k_retrieval)
            retrieved_ids = extract_retrieved_ids(raw_output)
                        
            # Use the new Hit Rate logic
            hit_score = calculate_hit_rate(retrieved_ids, gold_docs, sources_map)
            total_hit_rate += hit_score
            processed_count += 1
                        
            log_entry = {
                "question_id": idx + 1,
                "question": question,
                "gold_documents_raw": gold_docs,
                "retrieved_ids": retrieved_ids,
                "hit_score": hit_score
            }
            log_file.write(json.dumps(log_entry) + "\n")
                        
            if progress_callback: progress_callback("Gold Set Benchmark", processed_count, total_questions)
            time.sleep(3)

    average_hit_rate = total_hit_rate / processed_count if processed_count > 0 else 0.0
        
    summary_path = os.path.join(output_base_dir, "benchmark_summary.txt")
    with open(summary_path, "w", encoding="utf-8") as f:
        f.write("INFRAMAT-X RAG BENCHMARK REPORT\n")
        f.write(f"Run completed at: {timestamp}\n")
        f.write(f"Questions processed: {processed_count}\n")
        f.write(f"Average Hit Rate@10: {average_hit_rate:.4f}\n")
            
    summary_str = (
        f"✅ Benchmark finished!\n"
        f"📁 Logs saved to: {jsonl_path}\n"
        f"📊 Average Hit Rate@10: {average_hit_rate:.4f}\n"
    )
        
    import shutil
    zip_path = shutil.make_archive(output_base_dir, 'zip', output_base_dir)
        
    return summary_str, zip_path