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

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 ""

    href_match = re.search(r'href=["\'](?:https?://(?:huggingface\.co|hf\.co)/)?([^"\']+)["\']', name)
    if href_match:
        name = href_match.group(1)
    else:
        md_match = re.search(r'\((?:https?://(?:huggingface\.co|hf\.co)/)?([^)]+)\)', name)
        if md_match:
            name = md_match.group(1)

    name = re.sub(r'<[^>]+>', '', name)
    name = re.sub(r'\[([^\]]+)\]', r'\1', name)

    for domain in ["https://huggingface.co/", "http://huggingface.co/", "huggingface.co/",
                   "https://hf.co/", "http://hf.co/", "hf.co/"]:
        name = name.replace(domain, "")
    name = name.split("?")[0].split("#")[0]
    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 {}

    cert = mongo_get_certificate(clean_name)
    if cert and cert.get("status") == "ok":
        return cert

    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:
    if not cert or cert.get("status") != "ok":
        return ""

    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)

    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:
    """ZeroGPU execution worker - all parameters are strictly pickleable primitives."""
    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>ibm-granite/granite-4.1-3b-base</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."),
            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...", 8),
        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 = [
        (15, "Connecting to HF Hub & verifying repository metadata..."),
        (30, "Downloading weight tensors & initializing model configuration..."),
        (48, "Layer A: Performing SVD Spectral Tomography across weight matrices..."),
        (68, "Layer B: Streaming activations & fitting covariance observer..."),
        (84, "Layer C: Executing empirical factual probe battery..."),
        (94, "Finalizing operator risk analysis & generating certificate..."),
    ]

    start_time = time.time()
    while not done_event.is_set():
        elapsed = time.time() - start_time
        if elapsed < 3.0:
            pct, msg = stages[0]
        elif elapsed < 8.0:
            pct, msg = stages[1]
        elif elapsed < 18.0:
            pct, msg = stages[2]
        elif elapsed < 32.0:
            pct, msg = stages[3]
        elif elapsed < 48.0:
            pct, msg = stages[4]
        else:
            pct, msg = stages[5]

        yield (
            styled_loading(f"🔬 Auditing <b>{clean_model}</b> on ZeroGPU...", msg, pct),
            current_df,
            current_top3,
            current_panel,
        )
        done_event.wait(timeout=0.8)

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

    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,
        )