import json
import os
import re
import sys
import threading
import time
import traceback
from datetime import datetime, timezone

# Hugging Face ZeroGPU compatibility
try:
    import spaces
    has_spaces = True
except ImportError:
    has_spaces = False

def gpu_decorator(duration=120):
    def decorator(fn):
        if has_spaces and hasattr(spaces, "GPU"):
            return spaces.GPU(duration=duration)(fn)
        return fn
    return decorator

# Database integration
try:
    from src.db import mongo_get_certificate, mongo_save_certificate
except ImportError:
    from db import mongo_get_certificate, mongo_save_certificate

# Engine import
try:
    from src.audit.engine import AuditError, check_feasibility, run_full_audit
except ImportError:
    try:
        from src.engine import AuditError, check_feasibility, run_full_audit
    except ImportError:
        try:
            from audit.engine import AuditError, check_feasibility, run_full_audit
        except ImportError:
            from engine import AuditError, check_feasibility, run_full_audit

# Formatting import
try:
    from src.display.formatting import (
        build_top_3_cards_html,
        render_audit_details_panel,
        styled_error,
        styled_loading,
        styled_message,
    )
except ImportError:
    try:
        from src.formatting import (
            build_top_3_cards_html,
            render_audit_details_panel,
            styled_error,
            styled_loading,
            styled_message,
        )
    except ImportError:
        try:
            from display.formatting import (
                build_top_3_cards_html,
                render_audit_details_panel,
                styled_error,
                styled_loading,
                styled_message,
            )
        except ImportError:
            from formatting import (
                build_top_3_cards_html,
                render_audit_details_panel,
                styled_error,
                styled_loading,
                styled_message,
            )

# Utils import
try:
    from src.display.utils import BENCHMARK_COLS, COLS
except ImportError:
    try:
        from src.utils import BENCHMARK_COLS, COLS
    except ImportError:
        try:
            from display.utils import BENCHMARK_COLS, COLS
        except ImportError:
            from utils import BENCHMARK_COLS, COLS

# Envs & Populate imports
try:
    from src.envs import AUDIT_DEVICE, EVAL_RESULTS_PATH, MAX_AUDIT_PARAMS_BILLION, TOKEN
except ImportError:
    from envs import AUDIT_DEVICE, EVAL_RESULTS_PATH, MAX_AUDIT_PARAMS_BILLION, TOKEN

try:
    from src.populate import get_leaderboard_df, get_top_3_eval_cards
except ImportError:
    from populate import get_leaderboard_df, get_top_3_eval_cards


def clean_model_name(raw_name: str) -> str:
    """Robust extraction and sanitization of Hugging Face Model IDs."""
    if not raw_name:
        return ""
    name = str(raw_name).strip()
    if name.lower() in ("none", "null", "undefined", ""):
        return ""

    # Extract from href="..." or markdown [text](url)
    href_match = re.search(r'href=["\'](?:https?://huggingface\.co/)?([^"\']+)["\']', name)
    if href_match:
        name = href_match.group(1)
    else:
        md_match = re.search(r'\((?:https?://huggingface\.co/)?([^)]+)\)', name)
        if md_match:
            name = md_match.group(1)

    # Strip HTML tags & markdown
    name = re.sub(r'<[^>]+>', '', name)
    name = re.sub(r'\[([^\]]+)\]', r'\1', name)

    # Clean domain & query strings
    name = name.replace("https://huggingface.co/", "").replace("http://huggingface.co/", "")
    name = name.split("?")[0].split("#")[0]

    # Strip trailing branch paths like /tree/main, /blob/main
    name = re.sub(r'/(tree|blob|resolve)/.*$', '', name)

    return name.strip().strip("/")


def get_certificate_by_model_name(model_name: str) -> dict:
    clean_name = clean_model_name(model_name)
    if not clean_name:
        return {}

    # 1. Look up in MongoDB Atlas
    cert = mongo_get_certificate(clean_name)
    if cert and cert.get("status") == "ok":
        return cert

    # 2. Fallback to local files
    if os.path.exists(EVAL_RESULTS_PATH):
        safe_prefix = clean_name.replace("/", "__").lower()
        for f in os.listdir(EVAL_RESULTS_PATH):
            if f.lower().startswith(safe_prefix) and f.endswith(".json"):
                try:
                    with open(os.path.join(EVAL_RESULTS_PATH, f), "r", encoding="utf-8") as fp:
                        return json.load(fp)
                except Exception:
                    continue
    return {}


def _save_cert(model_id: str, revision: str, cert: dict) -> str:
    # Only save valid audit certificates to permanent storage
    if not cert or cert.get("status") != "ok":
        return ""

    # 1. Save to local disk cache
    os.makedirs(EVAL_RESULTS_PATH, exist_ok=True)
    safe_name = model_id.replace("/", "__") + f"_{revision}.json"
    out_path = os.path.join(EVAL_RESULTS_PATH, safe_name)
    try:
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(cert, f, indent=2)
    except Exception as e:
        print(f"Local file write error: {e}", flush=True)

    # 2. Persist permanently to MongoDB Atlas
    mongo_save_certificate(cert)

    return out_path


@gpu_decorator(duration=120)
def execute_direct_xray_audit(model_id: str, revision: str = "main", trust_remote_code: bool = True) -> dict:
    now_str = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")

    print(f"[LLM-X-RAY] Verifying repository metadata for '{model_id}'...", flush=True)

    try:
        feas = check_feasibility(model_id, revision, TOKEN, MAX_AUDIT_PARAMS_BILLION, trust_remote_code=trust_remote_code)
        if not feas.ok:
            return {
                "status": "error",
                "config": {"model_name": model_id, "model_sha": revision, "params": feas.param_count_b},
                "error_message": feas.reason,
                "audited_at": now_str,
            }

        result = run_full_audit(
            model_id=model_id,
            revision=revision,
            device="cuda" if has_spaces else AUDIT_DEVICE,
            token=TOKEN,
            trust_remote_code=trust_remote_code,
            progress_callback=None,
        )

        return {
            "status": "ok",
            "config": {
                "model_name": model_id,
                "model_sha": revision,
                "architecture": feas.architecture or "CausalLM",
                "params": feas.param_count_b or 0.5,
                "precision": "bfloat16",
                "license": "open-source",
            },
            "audited_at": now_str,
            **result,
        }
    except Exception as e:
        traceback.print_exc()
        sys.stdout.flush()
        return {
            "status": "error",
            "config": {"model_name": model_id, "model_sha": revision},
            "error_message": str(e),
            "audited_at": now_str,
        }


def audit_or_search_model(url_or_id: str, trust_remote_code: bool = True):
    clean_model = clean_model_name(url_or_id)

    current_df = get_leaderboard_df(EVAL_RESULTS_PATH, "", COLS, BENCHMARK_COLS)
    current_top3 = build_top_3_cards_html(get_top_3_eval_cards(EVAL_RESULTS_PATH))
    current_panel = render_audit_details_panel({})

    if not clean_model:
        yield (
            styled_error("Please enter a valid Hugging Face Model ID or URL (e.g. <code>Qwen/Qwen2.5-0.5B-Instruct</code>)."),
            current_df,
            current_top3,
            current_panel,
        )
        return

    if clean_model.startswith("spaces/"):
        yield (
            styled_error(f"<b>{clean_model}</b> is a Hugging Face <b>Space</b>, not a Model. Please enter a Model ID (e.g. <code>Qwen/Qwen2.5-0.5B-Instruct</code> or <code>SupraLabs/Supra2-Nano</code>)."),
            current_df,
            current_top3,
            current_panel,
        )
        return
    if clean_model.startswith("datasets/"):
        yield (
            styled_error(f"<b>{clean_model}</b> is a <b>Dataset</b>, not a Model. Please enter a Model ID."),
            current_df,
            current_top3,
            current_panel,
        )
        return

    # Check for existing cert in MongoDB Atlas / local cache (Instant Return)
    existing_cert = get_certificate_by_model_name(clean_model)
    if existing_cert and existing_cert.get("status") == "ok":
        yield (
            styled_message(f"Loaded existing audit certificate for <b>{clean_model}</b> from MongoDB Atlas."),
            current_df,
            current_top3,
            render_audit_details_panel(existing_cert),
        )
        return

    yield (
        styled_loading(f"Connecting to HF Hub for <b>{clean_model}</b>...", "Requesting ZeroGPU slice & downloading model weights..."),
        current_df,
        current_top3,
        current_panel,
    )

    result_holder = {}
    done_event = threading.Event()

    def _worker():
        try:
            result_holder["cert"] = execute_direct_xray_audit(
                clean_model, "main", trust_remote_code=trust_remote_code
            )
        except Exception as err:
            traceback.print_exc()
            sys.stdout.flush()
            result_holder["cert"] = {
                "status": "error",
                "config": {"model_name": clean_model},
                "error_message": str(err),
            }
        finally:
            done_event.set()

    thread = threading.Thread(target=_worker, daemon=True)
    thread.start()

    stages = [
        "Downloading model weights and tokenizer tensors...",
        "Layer A: Performing SVD Spectral Tomography across weight tensors...",
        "Layer B: Streaming activations and calculating operator covariance...",
        "Layer C: Running empirical factual probe battery...",
        "Finalizing operator risk and registering certificate...",
    ]
    stage_idx = 0

    while not done_event.is_set():
        current_stage = stages[min(stage_idx, len(stages) - 1)]
        yield (
            styled_loading(f"🔬 Auditing <b>{clean_model}</b> on ZeroGPU...", current_stage),
            current_df,
            current_top3,
            current_panel,
        )
        done_event.wait(timeout=2.5)
        stage_idx += 1

    thread.join()
    cert = result_holder.get("cert", {})

    # Save to MongoDB and local disk only on success
    if cert and cert.get("status") == "ok":
        _save_cert(clean_model, cert.get("config", {}).get("model_sha", "main"), cert)

    updated_df = get_leaderboard_df(EVAL_RESULTS_PATH, "", COLS, BENCHMARK_COLS)
    updated_top3 = build_top_3_cards_html(get_top_3_eval_cards(EVAL_RESULTS_PATH))
    updated_panel = render_audit_details_panel(cert)

    if cert.get("status") == "error":
        err_msg = cert.get("error_message", "Unknown error.")
        yield (
            styled_error(f"Audit could not be completed for <b>{clean_model}</b>: {err_msg}"),
            updated_df,
            updated_top3,
            updated_panel,
        )
    else:
        yield (
            styled_message(f"🎉 <b>Audit Complete for {clean_model}!</b> Persisted in MongoDB Atlas."),
            updated_df,
            updated_top3,
            updated_panel,
        )