"""LoRA fetch + load from Hugging Face (or any direct .safetensors path).

Functions:
- `_parse_hf_link(url)`            → `(repo, file | None)`
- `_resolve_repo_file(repo)`       → pick the single .safetensors of a repo link
- `_resolve_repo_revision(repo)`   → Hub tip SHA (immutable cache key)
- `_lora_cache_path(repo, file, revision)` → bucket path, or None if not writable
- `_fetch_lora_file(repo, file, revision)` → local .safetensors path (cached)
- `_load_lora_from_cache(path)`    → load any local .safetensors into _CUSTOM_LoRA
- `_ensure_nsfw_lora_loaded(key)`  → lazy fetch + convert a catalog NSFW LoRA
- `_ensure_custom_lora_loaded(repo, file)` → same for a user-supplied HF LoRA
- `_ensure_custom_lora_loaded_from_path(key, path, label)` → load an already-fetched file

Three layers of `H3_MAX_LORA_BYTES` enforcement: HEAD pre-check, post-download
`os.path.getsize`, and a running byte counter during the bucket-cache copy
(defends against a Hub lying about Content-Length). Concurrent writers use
`_LORA_CACHE_LOCK` + temp file + `os.replace`.
"""
from __future__ import annotations

import json
import os
import re
import struct
import time

import gradio as gr

from config import NSFW_LORA_CONFIGS
from lora.cache import _evict_lora_cache_for_quota
from lora.formats import (
    _convert_lora,
    _detect_wrong_arch,
)
import state as _state
from state import (
    _CUSTOM_LoRA,
    _H3_MAX_LORA_BYTES,
    _LORA_CACHE_DIR,
    _LORA_CACHE_LOCK,
    _NSFW_LoRA,
    _log,
)


def _parse_hf_link(url):
    """Parse a Hugging Face repo or file link into (repo, file). `file` is None for a bare repo link."""
    url = (url or "").strip().split("?")[0].split("#")[0]
    if not url:
        raise gr.Error("Paste a Hugging Face or CivitAI link first.")
    if not url.startswith("http"):
        parts = [p for p in url.split("/") if p]
        if len(parts) >= 2:
            return parts[0] + "/" + parts[1], "/".join(parts[2:]) or None
        raise gr.Error(f"Couldn't parse `{url}` as a Hugging Face repo link.")
    if "huggingface.co" not in url:
        raise gr.Error("Only Hugging Face links are supported (huggingface.co).")
    parts = [p for p in url.split("huggingface.co", 1)[1].strip("/").split("/") if p]
    if parts and parts[0] in ("spaces", "models", "datasets"):
        parts = parts[1:]  # https://huggingface.co/spaces/Owner/Repo → Owner/Repo
    if len(parts) < 2:
        raise gr.Error(f"Couldn't parse `{url}` as a Hugging Face repo link.")
    repo = parts[0] + "/" + parts[1]
    if len(parts) == 2:
        return repo, None
    if parts[2] in ("resolve", "tree", "blob"):
        if len(parts) == 4 and parts[2] in ("tree", "blob"):
            return repo, None  # tree/main or blob/main with no subpath → treat as a repo link
        if len(parts) < 5:
            raise gr.Error(f"Couldn't parse a filename from `{url}`.")
        if parts[3] != "main":
            raise gr.Error("Only `main`-branch links are supported — use a `resolve/main/...` link.")
        return repo, "/".join(parts[4:])
    return repo, "/".join(parts[2:])


def _resolve_repo_file(repo):
    """Pick the single .safetensors file of a repo link; error if there isn't exactly one."""
    from huggingface_hub import list_repo_files
    try:
        files = [f for f in list_repo_files(repo) if f.endswith(".safetensors")]
    except Exception as _e:
        raise gr.Error(f"Couldn't list files in `{repo}`: {_e}")
    if len(files) == 1:
        return files[0]
    if not files:
        raise gr.Error(f"No .safetensors files found in `{repo}`.")
    raise gr.Error(f"`{repo}` has multiple .safetensors files — paste a direct file link instead: "
                   + ", ".join(files[:5]) + ("…" if len(files) > 5 else ""))


def _resolve_repo_revision(repo: str) -> str:
    """Resolve the current tip of *repo* to an immutable commit SHA.

    Falls back to the literal string 'main' only if the Hub is unreachable so
    cache keys remain stable within a process; callers must not treat 'main' as
    immutable across restarts.
    """
    try:
        from huggingface_hub import HfApi
        info = HfApi().repo_info(repo, repo_type="model")
        sha = getattr(info, "sha", None) or getattr(info, "revision", None)
        if sha and isinstance(sha, str) and len(sha) >= 7:
            return sha
    except Exception as _e:
        _log(f"revision resolve failed for {repo}: {_e}", "WARN")
    return "main"


def _lora_cache_path(repo, file, revision: str | None = None):
    """Bucket path for a (repo, revision, file) LoRA, or None if the bucket isn't usable.

    Cache identity is (repo, revision_sha, file) so a moving `main` cannot serve
    stale weights under the same path.
    """
    try:
        os.makedirs(_LORA_CACHE_DIR, exist_ok=True)
        _probe = os.path.join(_LORA_CACHE_DIR, ".probe")
        with open(_probe, "w") as _f:
            _f.write("ok")
        os.remove(_probe)
    except Exception:
        return None
    rev = (revision or "main")[:40]
    _key = re.sub(r"[^A-Za-z0-9._-]", "_", f"{repo.replace('/', '__')}__rev_{rev}__{file}")
    return os.path.join(_LORA_CACHE_DIR, _key)


def _fetch_lora_file(repo, file, revision: str | None = None, max_bytes: int | None = None):
    """Local path to a LoRA .safetensors, preferring the bucket cache over a fresh download.

    Cache key includes the immutable Hub revision SHA when available.
    Enforces H3_MAX_LORA_BYTES (reject before/during download) and
    H3_MAX_LORA_CACHE_BYTES (evict unreferenced entries or reject).
    Concurrent writers use temp files + os.replace so consumers never see
    partial artifacts.

    ``max_bytes`` overrides the user-LoRA size cap (used for the system turbo
    weights, which are larger than the custom-LoRA quota by design).
    """
    from huggingface_hub import hf_hub_download, hf_hub_url
    import requests

    limit = int(max_bytes) if max_bytes is not None else _H3_MAX_LORA_BYTES
    rev = revision or _resolve_repo_revision(repo)
    with _LORA_CACHE_LOCK:
        _cached = _lora_cache_path(repo, file, revision=rev)
        if _cached and os.path.exists(_cached):
            try:
                with open(_cached, "rb") as _f:
                    _n = struct.unpack("<Q", _f.read(8))[0]
                    _header_keys = list(json.loads(_f.read(_n)).keys())
                _hint = _detect_wrong_arch(_header_keys)
                if _hint:
                    _log(f"LoRA cache evict (incompatible): {repo}/{file}@{rev[:7]} — {_hint}", "WARN")
                    try:
                        os.remove(_cached)
                    except OSError:
                        pass
                    raise gr.Error(f"`{file}` is not a MiniMax-H3 LoRA. {_hint}")
                size = os.path.getsize(_cached)
                if size > limit:
                    try:
                        os.remove(_cached)
                    except OSError:
                        pass
                    raise gr.Error(
                        f"`{file}` exceeds H3_MAX_LORA_BYTES "
                        f"({size / 1e6:.0f} MB > {limit / 1e6:.0f} MB)."
                    )
                _log(f"LoRA cache hit: {repo}/{file}@{rev[:7]} ({size / 1e6:.0f} MB)")
                return _cached
            except gr.Error:
                raise
            except Exception:
                _log(f"LoRA cache corrupt, re-downloading: {repo}/{file}@{rev[:7]}", "WARN")
                try:
                    os.remove(_cached)
                except OSError:
                    pass

        # Pre-check size via HEAD when Content-Length is available.
        try:
            url = hf_hub_url(repo_id=repo, filename=file, revision=rev if rev != "main" else "main")
            head = requests.head(url, allow_redirects=True, timeout=30)
            cl = head.headers.get("Content-Length")
            if cl is not None:
                expected = int(cl)
                if expected > limit:
                    raise gr.Error(
                        f"`{file}` is too large ({expected / 1e6:.0f} MB > "
                        f"{limit / 1e6:.0f} MB limit)."
                    )
                _evict_lora_cache_for_quota(expected)
        except gr.Error:
            raise
        except Exception as _e:
            _log(f"LoRA HEAD pre-check skipped ({_e})", "WARN")

        _t0 = time.time()
        # Download via hub; then enforce size on the resulting file before cache admit.
        _path = hf_hub_download(repo, file, revision=rev if rev != "main" else None)
        size = os.path.getsize(_path)
        _log(f"LoRA downloaded: {repo}/{file}@{rev[:7]} ({size / 1e6:.0f} MB in {time.time() - _t0:.0f}s)")
        if size > limit:
            try:
                os.remove(_path)
            except OSError:
                pass
            raise gr.Error(
                f"`{file}` exceeds H3_MAX_LORA_BYTES "
                f"({size / 1e6:.0f} MB > {limit / 1e6:.0f} MB)."
            )

        # Architecture gate before admitting to persistent cache.
        try:
            with open(_path, "rb") as _f:
                _n = struct.unpack("<Q", _f.read(8))[0]
                _header = json.loads(_f.read(_n))
            _hint = _detect_wrong_arch(list(_header.keys()))
            if _hint:
                _log(f"LoRA rejected (incompatible architecture): {repo}/{file} — {_hint}", "WARN")
                try:
                    os.remove(_path)
                except OSError:
                    pass
                raise gr.Error(f"`{file}` is not a MiniMax-H3 LoRA. {_hint}")
        except gr.Error:
            raise
        except Exception as _e:
            _log(f"LoRA header peek failed ({_e}); proceeding with copy", "WARN")

        if _cached:
            try:
                _evict_lora_cache_for_quota(size)
                _tmp = _cached + ".tmp"
                # Stream with a running byte counter so a lied Content-Length still cannot exceed the ceiling.
                written = 0
                with open(_path, "rb") as _src, open(_tmp, "wb") as _dst:
                    while True:
                        chunk = _src.read(1024 * 1024)
                        if not chunk:
                            break
                        written += len(chunk)
                        if written > limit:
                            _dst.close()
                            try:
                                os.remove(_tmp)
                            except OSError:
                                pass
                            raise gr.Error(
                                f"`{file}` exceeded H3_MAX_LORA_BYTES while copying "
                                f"({written / 1e6:.0f} MB)."
                            )
                        _dst.write(chunk)
                os.replace(_tmp, _cached)
                _log(f"LoRA cached to bucket: {repo}/{file}@{rev[:7]}")
                return _cached
            except gr.Error:
                raise
            except Exception as _e:
                _log(f"LoRA bucket cache write failed ({_e}); using HF cache", "WARN")
                try:
                    if os.path.exists(_cached + ".tmp"):
                        os.remove(_cached + ".tmp")
                except OSError:
                    pass
        return _path


def _ensure_nsfw_lora_loaded(key):
    """Download + convert + store a lazy NSFW LoRA's A/B factors on CPU. No-op if already loaded."""
    if key in _NSFW_LoRA:
        return
    cfg = NSFW_LORA_CONFIGS[key]
    import torch
    from safetensors.torch import load_file as _load_lora
    _t0 = time.time()
    _log(f"loading NSFW LoRA {key} from {cfg['repo']} ...")
    # Catalog NSFW weights can exceed the user custom-LoRA quota (default 512 MB);
    # allow up to 1.5 GB for official catalog entries only.
    _path = _fetch_lora_file(cfg["repo"], cfg["file"], max_bytes=1536 * 1024 * 1024)
    _raw = _load_lora(_path)
    # The catalog's `format` is a hint, not a contract — a mislabeled entry converts under whatever
    # format its keys actually are, and only a file no converter recognizes becomes a user error.
    _sd, _fmt = _convert_lora(_raw, cfg.get("format", ""))
    if not _sd:
        raise gr.Error(
            f"`{cfg['file']}` converted to no usable modules — it isn't a MiniMax-H3 LoRA "
            f"(or its `format` is mislabeled). {_detect_wrong_arch(_raw)}"
        )
    _log(f"NSFW LoRA {key}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(
            f"`{cfg['file']}` has no complete LoRA module pairs — it isn't a MiniMax-H3 LoRA "
            f"(or its `format` is mislabeled). {_detect_wrong_arch(_raw)}"
        )
    del _sd, _pairs, _raw
    import gc as _gc; _gc.collect()
    _NSFW_LoRA[key] = _loaded
    _log(f"NSFW LoRA {key} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")


def _ensure_custom_lora_loaded(repo, file):
    """Download + convert + store a user-supplied Hugging Face LoRA's A/B factors on CPU. No-op if already loaded."""
    key = (repo, file)
    if key in _CUSTOM_LoRA:
        return
    import torch
    from safetensors.torch import load_file as _load_lora
    _t0 = time.time()
    _log(f"loading custom LoRA {repo}/{file} ...")
    try:
        _path = _fetch_lora_file(repo, file)
        _raw = _load_lora(_path)
    except Exception as _e:
        raise gr.Error(f"Couldn't fetch custom LoRA `{repo}/{file}`: {_e}")
    _sd, _fmt = _convert_lora(_raw)
    if not _sd:
        raise gr.Error(f"`{repo}/{file}` doesn't look like a MiniMax-H3 LoRA (no lora_A/lora_down/lokr keys found).")
    _log(f"custom LoRA {repo}/{file}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(f"`{repo}/{file}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
    del _sd, _pairs, _raw
    import gc as _gc; _gc.collect()
    _CUSTOM_LoRA[key] = _loaded
    _log(f"custom LoRA {repo}/{file} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")


def _load_lora_from_cache(path):
    """Load a cached .safetensors LoRA directly by file path.  Returns (repo, file) key
    and populates ``_CUSTOM_LoRA`` so it can be merged on generate."""
    import torch
    from safetensors.torch import load_file as _load_sd
    _t0 = time.time()
    fname = os.path.basename(path)
    # Use the filename as a synthetic (repo, file) key.
    key = ("__cache__", fname)
    if key in _CUSTOM_LoRA:
        return key
    # Architecture gate: peek at header keys before loading full tensor data.
    try:
        with open(path, "rb") as _f:
            _n = struct.unpack("<Q", _f.read(8))[0]
            _header_keys = list(json.loads(_f.read(_n)).keys())
        _hint = _detect_wrong_arch(_header_keys)
        if _hint:
            _log(f"cached LoRA rejected (incompatible): {fname} — {_hint}", "WARN")
            try:
                os.remove(path)
            except OSError:
                pass
            raise gr.Error(f"`{fname}` is not a MiniMax-H3 LoRA. {_hint}")
    except gr.Error:
        raise
    except Exception:
        pass  # header peek failed — fall through to full load
    try:
        _raw = _load_sd(path)
    except Exception as _e:
        raise gr.Error(f"Couldn't load cached LoRA `{fname}`: {_e}")
    # Detect format and convert.
    _sd, _fmt = _convert_lora(_raw)
    if not _sd:
        raise gr.Error(f"`{fname}` doesn't look like a MiniMax-H3 LoRA.")
    _log(f"cached LoRA {fname}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(f"`{fname}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
    del _sd, _pairs, _raw
    import gc as _gc; _gc.collect()
    _CUSTOM_LoRA[key] = _loaded
    _log(f"cached LoRA {fname} loaded ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")
    return key


def _ensure_custom_lora_loaded_from_path(key, path: str, label: str = ""):
    """Load + convert a local .safetensors LoRA into ``_CUSTOM_LoRA[key]``. No-op if present.

    Used by the CivitAI / CivArchive flows after they've fetched the file.
    """
    if key in _CUSTOM_LoRA:
        return
    import torch
    from safetensors.torch import load_file as _load_lora

    _t0 = time.time()
    label = label or f"{key[0]}/{key[1]}"
    _log(f"loading custom LoRA {label} from local path ...")
    try:
        _raw = _load_lora(path)
    except Exception as _e:
        raise gr.Error(f"Couldn't read LoRA `{label}`: {_e}")
    _sd, _fmt = _convert_lora(_raw)
    if not _sd:
        raise gr.Error(f"`{label}` doesn't look like a MiniMax-H3 LoRA (no lora_A/lora_down/lokr keys found).")
    _log(f"custom LoRA {label}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(f"`{label}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
    del _sd, _pairs, _raw
    import gc as _gc
    _gc.collect()
    _CUSTOM_LoRA[key] = _loaded
    _log(f"custom LoRA {label} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")