"""Private, content-addressed MOSS transformer LoRAs. No executable model code.

Only standard A/B (up/down) linear LoRA on the 42 MOSS DiT blocks is supported.
Weights merge from immutable host originals; no merged checkpoint is written.
"""
from __future__ import annotations
import hashlib
import json
import math
import os
from pathlib import Path
import re
import threading

MAX_BYTES = 512 * 1024 * 1024
MAX_TOTAL_BYTES = 5 * 1024 * 1024 * 1024
MAX_MODULES = 4
ID = re.compile(r"^[a-f0-9]{64}$")
STORE_LOCK = threading.RLock()


def owner_hash(token):
    if not isinstance(token, str) or not re.fullmatch(r"[A-Za-z0-9_-]{43,128}", token):
        raise ValueError("A private LoRA client token is required")
    return hashlib.sha256(token.encode()).hexdigest()


def normalize_config(raw):
    config = json.loads(raw)
    if not isinstance(config, dict) or set(config) != {"alpha", "base_model"}:
        raise ValueError("Invalid LoRA configuration")
    alpha = config["alpha"]
    if alpha is not None and (isinstance(alpha, bool) or not isinstance(alpha, (int, float)) or not math.isfinite(alpha) or not 0 < alpha <= 4096):
        raise ValueError("Network alpha must be positive and at most 4096")
    model = config["base_model"]
    if model is not None and (not isinstance(model, str) or len(model) > 200):
        raise ValueError("Invalid adapter base model")
    return config


def target_key(name):
    name = re.sub(r"^(?:base_model\.model\.|transformer\.|diffusion_model\.)", "", name)
    name = re.sub(r"^lora_unet_layers_(\d+)_proj_(q|k|v|o|ffn_up_gate|ffn_down)$", r"layers.\1.proj.\2", name)
    match = re.fullmatch(r"layers\.(\d+)\.proj\.(q|k|v|o|ffn_up_gate|ffn_down)", name)
    if not match or int(match[1]) >= 42:
        raise ValueError("Unsupported tensor target: use a MOSS DiT linear LoRA (not SDXL, FLUX, LoKr or DoRA)")
    return name + ".weight", match[2]


def inspect_adapter(file, config):
    from safetensors import SafetensorError
    try:
        return _inspect_adapter(file, config)
    except SafetensorError as error:
        raise ValueError("Invalid or incomplete safetensors adapter") from error


def _inspect_adapter(file, config):
    from safetensors import safe_open
    pairs = {}
    with safe_open(str(file), framework="pt", device="cpu") as sf:
        metadata = sf.metadata() or {}
        for key in sf.keys():
            match = re.fullmatch(r"(.+)\.(lora_A(?:\.default)?\.weight|lora_B(?:\.default)?\.weight|lora_down\.weight|lora_up\.weight|alpha)", key)
            if not match:
                raise ValueError("Unsupported adapter tensors; only standard linear LoRA is supported")
            target, kind = target_key(match[1])
            group = pairs.setdefault(target, {"kind": kind})
            suffix = match[2]
            role = "alpha" if suffix == "alpha" else "a" if suffix.startswith(("lora_A", "lora_down")) else "b"
            if role in group:
                raise ValueError("Duplicate LoRA tensor target")
            shape = list(sf.get_slice(key).get_shape())
            if role == "alpha":
                if shape not in ([], [1]):
                    raise ValueError("Invalid per-layer network alpha")
                alpha = float(sf.get_tensor(key).item())
                if not math.isfinite(alpha) or not 0 < alpha <= 4096:
                    raise ValueError("Invalid per-layer network alpha")
                group[role] = alpha
            else:
                if sf.get_slice(key).get_dtype() not in ("F16", "BF16", "F32") or len(shape) != 2:
                    raise ValueError("LoRA matrices must be float16, bfloat16 or float32")
                group[role] = key
                group[role + "_shape"] = shape
        if not pairs:
            raise ValueError("No LoRA matrices found")
        for target, group in pairs.items():
            if "a" not in group or "b" not in group:
                raise ValueError("Incomplete LoRA matrix pair")
            a, b = group["a_shape"], group["b_shape"]
            expected = [16384, 3072] if group["kind"] == "ffn_up_gate" else [3072, 8192] if group["kind"] == "ffn_down" else [3072, 3072]
            if a[0] != b[1] or not 1 <= a[0] <= 512 or [b[0], a[1]] != expected:
                raise ValueError("Adapter tensor dimensions do not match MOSS-Image")
            alpha = group.get("alpha", config.get("alpha"))
            if alpha is None and metadata.get("ss_network_alpha") is not None:
                alpha = float(metadata["ss_network_alpha"])
            if alpha is None:
                alpha = a[0]
            if not math.isfinite(alpha) or not 0 < alpha <= 4096:
                raise ValueError("Invalid network alpha")
            group["factor"] = alpha / a[0]
    return pairs


def validate_selection(selection):
    if not isinstance(selection, list) or len(selection) > MAX_MODULES:
        raise ValueError("Choose at most four LoRA modules")
    seen = set()
    for item in selection:
        if not isinstance(item, dict) or set(item) != {"id", "strength"} or not ID.fullmatch(str(item.get("id", ""))):
            raise ValueError("Invalid LoRA module selection")
        strength = item["strength"]
        if isinstance(strength, bool) or not isinstance(strength, (int, float)) or not math.isfinite(strength) or not -2 <= strength <= 2:
            raise ValueError("LoRA strength must be between -2 and 2")
        if item["id"] in seen:
            raise ValueError("Duplicate LoRA module")
        seen.add(item["id"])
    return selection


class AdapterStore:
    def __init__(self, root, model):
        self.root = Path(root)
        self.model = model

    def paths(self, owner, identifier):
        if not ID.fullmatch(owner) or not ID.fullmatch(identifier):
            raise ValueError("Invalid LoRA identifier")
        return self.root / owner / (identifier + ".safetensors"), self.root / owner / (identifier + ".json")

    def compatible(self, config):
        declared = config.get("base_model")
        # Base-trained style adapters also support the compatible SFT model.
        if declared and declared != self.model and not (declared == "OpenMOSS-Team/MOSS-Image-Base" and self.model == "OpenMOSS-Team/MOSS-Image-SFT"):
            raise ValueError("Adapter declares a different base model from the serving model")

    def install(self, temporary, owner, identifier, config_raw, name):
        config = normalize_config(config_raw)
        self.compatible(config)
        digest = hashlib.sha256()
        with open(temporary, "rb") as handle:
            while chunk := handle.read(1024 * 1024):
                digest.update(chunk)
        calculated = hashlib.sha256((digest.hexdigest() + "\n" + config_raw).encode()).hexdigest()
        if calculated != identifier:
            raise ValueError("LoRA checksum mismatch")
        pairs = inspect_adapter(temporary, config)
        file, manifest = self.paths(owner, identifier)
        record = {"id": identifier, "name": name[:160], "bytes": Path(temporary).stat().st_size,
                  "sha256": digest.hexdigest(), "config": config, "config_raw": config_raw, "targets": len(pairs)}
        with STORE_LOCK:
            current = sum(p.stat().st_size for p in self.root.glob("*/*.safetensors"))
            if not file.exists() and current + record["bytes"] > MAX_TOTAL_BYTES:
                raise ValueError("Server LoRA cache is full; remove unused server modules")
            file.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
            os.replace(temporary, file)
            file.chmod(0o600)
            tmp = manifest.with_suffix(".json.tmp")
            with open(tmp, "w") as handle:
                json.dump(record, handle)
                handle.flush()
                os.fsync(handle.fileno())
            tmp.chmod(0o600)
            os.replace(tmp, manifest)
        return record

    def get(self, owner, identifier):
        file, manifest = self.paths(owner, identifier)
        if not file.is_file() or not manifest.is_file():
            raise FileNotFoundError("LoRA module not uploaded")
        record = json.loads(manifest.read_text())
        self.compatible(record["config"])
        return record

    def resolve(self, owner, selection):
        validate_selection(selection)
        resolved = []
        for item in selection:
            record = self.get(owner, item["id"])
            file, _ = self.paths(owner, item["id"])
            resolved.append({**record, "file": str(file), "strength": item["strength"]})
        return resolved


def merge_weights(base, adapters):
    import torch
    from safetensors import safe_open
    merged = dict(base)
    # Accumulate each target in fp32 and round once, always from immutable base.
    targets = {}
    for item in adapters:
        if item["strength"] == 0:
            continue
        digest = hashlib.sha256()
        with open(item["file"], "rb") as handle:
            while chunk := handle.read(1024 * 1024):
                digest.update(chunk)
        if digest.hexdigest() != item["sha256"]:
            raise ValueError("Cached LoRA checksum mismatch")
        for target, pair in inspect_adapter(item["file"], item["config"]).items():
            if target not in base:
                raise ValueError("Adapter target is absent from the loaded model")
            targets.setdefault(target, []).append((item, pair))
    for target, terms in targets.items():
        weight = base[target].float().clone()
        for item, pair in terms:
            with safe_open(item["file"], framework="pt", device="cpu") as sf:
                a, b = sf.get_tensor(pair["a"]).float(), sf.get_tensor(pair["b"]).float()
            if not torch.isfinite(a).all() or not torch.isfinite(b).all():
                raise ValueError("LoRA contains non-finite weights")
            weight.addmm_(b, a, beta=1, alpha=item["strength"] * pair["factor"])
        if not torch.isfinite(weight).all():
            raise ValueError("LoRA produces non-finite model weights")
        merged[target] = weight.to(base[target].dtype)
    return merged


def apply_adapters(pipe, adapters):
    fingerprint = tuple((item["id"], item["strength"]) for item in adapters if item["strength"] != 0)
    if getattr(pipe, "_lora_fingerprint", ()) == fingerprint:
        return
    if not hasattr(pipe, "_lora_original_weights"):
        pipe._lora_original_weights = pipe._load_dit_weights()
    merged = merge_weights(pipe._lora_original_weights, adapters)
    if pipe._dit_model is not None:
        from ttnn_moss.pipeline import _free_ttnn_module
        _free_ttnn_module(pipe._dit_model)
        pipe._dit_model = None
    pipe._dit_weights = merged
    pipe._lora_fingerprint = fingerprint