import os
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HOME"] = "/tmp/hf_cache"

import json, time, warnings, threading
from pathlib import Path
warnings.filterwarnings("ignore")

import gradio as gr

result_text = "Initializing..."

def run_task():
    global result_text
    try:
        import mteb
        import datasets
        import pandas as pd
        import pyarrow.parquet as pq
        
        MODEL = "mixedbread-ai/deepset-mxbai-embed-de-large-v1"
        
        cpuinfo = Path("/proc/cpuinfo").read_text()
        has_avx2 = "avx2" in cpuinfo.lower()
        
        benchmark = None
        for b in mteb.get_benchmarks():
            if getattr(b, "name", None) == "MTEB(eng, v2)":
                benchmark = b
                break
        
        mind_task = [t for t in benchmark.tasks if t.metadata.name == "MindSmallReranking"][0]
        result_text = f"AVX2: {has_avx2}\nTask: {mind_task.metadata.name}"
        
        # Download from our HF bucket (fast HF-to-HF)
        from huggingface_hub import snapshot_download
        result_text += "\nDownloading dataset..."
        t0 = time.time()
        snapshot_download(
            repo_id="mteb/MindSmallReranking",
            repo_type="dataset",
            revision="227478e3235572039f4f7661840e059f31ef6eb1",
        )
        result_text += f"\nDownloaded: {time.time()-t0:.1f}s"
        
        # Now load parquet files DIRECTLY using pyarrow (much faster than datasets.load_dataset)
        from huggingface_hub import hf_hub_download
        cache_dir = Path("/tmp/hf_cache/hub")
        ds_repo = "datasets--mteb--MindSmallReranking"
        snap_dir = list(cache_dir.glob(f"{ds_repo}/snapshots/*"))[0]
        
        result_text += "\nLoading parquet files directly..."
        t0 = time.time()
        
        # Load qrels (default config = data/ folder)
        qrels_files = sorted((snap_dir / "data").glob("test-*.parquet"))
        result_text += f"\nQrels files: {len(qrels_files)}"
        
        # Load queries
        queries_files = sorted((snap_dir / "queries").glob("test-*.parquet"))
        queries_df = pd.concat([pd.read_parquet(f) for f in queries_files], ignore_index=True)
        result_text += f"\nQueries: {len(queries_df)} rows"
        
        # Load corpus
        corpus_files = sorted((snap_dir / "corpus").glob("test-*.parquet"))
        corpus_df = pd.concat([pd.read_parquet(f) for f in corpus_files], ignore_index=True)
        result_text += f"\nCorpus: {len(corpus_df)} rows"
        
        # Load top_ranked (for reranking)
        tr_files = sorted((snap_dir / "top_ranked").glob("test-*.parquet"))
        top_ranked_df = pd.concat([pd.read_parquet(f) for f in tr_files], ignore_index=True)
        result_text += f"\nTop-ranked: {len(top_ranked_df)} rows"
        
        # Load qrels from data/
        qrels_df = pd.concat([pd.read_parquet(f) for f in qrels_files], ignore_index=True)
        result_text += f"\nQrels: {len(qrels_df)} rows"
        
        elapsed = time.time() - t0
        result_text += f"\nAll parquet loaded: {elapsed:.1f}s"
        result_text += f"\nColumns: qrels={list(qrels_df.columns)}, queries={list(queries_df.columns)}, corpus={list(corpus_df.columns)}, top_ranked={list(top_ranked_df.columns)}"
        
        # Build the data structures that MTEB expects
        # qrels: {query_id: {corpus_id: score}}
        qrels_dict = {}
        for _, row in qrels_df.iterrows():
            qid = str(row["query-id"])
            cid = str(row["corpus-id"])
            score = int(row["score"])
            if qid not in qrels_dict:
                qrels_dict[qid] = {}
            qrels_dict[qid][cid] = score
        result_text += f"\nQrels dict: {len(qrels_dict)} queries"
        
        # queries: Dataset with id and text columns
        queries_ds = datasets.Dataset.from_pandas(queries_df[["id", "text"]])
        result_text += f"\nQueries dataset: {len(queries_ds)}"
        
        # corpus: Dataset with id and text columns  
        corpus_ds = datasets.Dataset.from_pandas(corpus_df[["id", "text"]])
        result_text += f"\nCorpus dataset: {len(corpus_ds)}"
        
        # top_ranked: {query_id: [corpus_id, ...]}
        top_ranked_dict = {}
        for _, row in top_ranked_df.iterrows():
            qid = str(row["query-id"])
            cids = [str(c) for c in row["corpus-ids"]]
            top_ranked_dict[qid] = cids
        result_text += f"\nTop-ranked dict: {len(top_ranked_dict)} queries"
        
        # Inject data into the task object (plain dict, not RetrievalSplitData)
        mind_task.dataset = {
            "default": {
                "test": {
                    "corpus": corpus_ds,
                    "queries": queries_ds,
                    "relevant_docs": qrels_dict,
                    "top_ranked": top_ranked_dict,
                }
            }
        }
        mind_task.data_loaded = True
        result_text += "\nData injected into task!"
        
        # Load model
        result_text += "\nLoading model..."
        t0 = time.time()
        model = mteb.get_model(MODEL)
        result_text += f"\nModel loaded: {time.time()-t0:.1f}s"
        
        # Run evaluation
        result_text += "\nRunning MindSmallReranking..."
        outdir = Path("/tmp/output")
        outdir.mkdir(parents=True, exist_ok=True)
        
        t0 = time.time()
        results = mteb.evaluate(
            model, tasks=[mind_task],
            prediction_folder=str(outdir),
            overwrite_strategy="always", raise_error=True,
        )
        elapsed = time.time() - t0
        result_text += f"\nCompleted: {elapsed:.0f}s ({elapsed/60:.1f} min)\n"
        
        for tr in results.task_results:
            for split, sv in tr.scores.items():
                if isinstance(sv, list):
                    for s in sv:
                        ms = s.get("main_score")
                        if ms is not None:
                            result_text += f"SCORE: {tr.task_name} [{split}]: {ms:.4f}\n"
                elif isinstance(sv, dict):
                    ms = sv.get("main_score")
                    if ms is not None:
                        result_text += f"SCORE: {tr.task_name} [{split}]: {ms:.4f}\n"
        
        from mteb.results.task_result import TaskResult
        tr_data = results.task_results[0]
        task_result = TaskResult.model_validate(tr_data.model_dump())
        result_text += "=== JSON_START ===\n"
        result_text += task_result.model_dump_json(indent=2)
        result_text += "\n=== JSON_END ==="
        
    except Exception as e:
        import traceback
        result_text = f"ERROR: {e}\n{traceback.format_exc()}"

def get_status():
    return result_text

threading.Thread(target=run_task, daemon=True).start()

with gr.Blocks() as demo:
    gr.Markdown("# MTEB MindSmallReranking Runner")
    out = gr.TextArea(label="Status", value=result_text, lines=30)
    timer = gr.Timer(value=10)
    timer.tick(get_status, outputs=out)

demo.launch(server_name="0.0.0.0", server_port=7860)