#!/usr/bin/env python3
"""HF MeshAI-Base-Models (dataset) uzerinden TRELLIS/Hunyuan agirlik indirme + LoRA hedef secimi."""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn

BASE_DATASET_REPO = "HayrettinIscan/MeshAI-Base-Models"

# A100 icin makul ilk paket (Hunyuan DiT ~5GB — VM disk yeterli)
DEFAULT_WEIGHT_FILES = (
    "Microsoft-TRELLIS/ckpts/slat_enc_swin8_B_64l8_fp16.safetensors",
    "Microsoft-TRELLIS/ckpts/ss_enc_conv3d_16l8_fp16.safetensors",
    "Microsoft-TRELLIS/ckpts/slat_dec_mesh_swin8_B_64l8m256c_fp16.safetensors",
    "Tencent-Hunyuan3D/hunyuan3d-dit-v2-0-turbo/model.fp16.safetensors",
)


def ensure_base_weights(
    *,
    token: str,
    cache_dir: Path,
    files: tuple[str, ...] = DEFAULT_WEIGHT_FILES,
    log_fn: Any = print,
) -> dict[str, Path]:
    """Indirilen safetensors yollarini dondurur (HF hub cache veya local_dir)."""
    from huggingface_hub import hf_hub_download

    cache_dir.mkdir(parents=True, exist_ok=True)
    out: dict[str, Path] = {}
    for rel in files:
        log_fn(f"[faz2] indiriliyor: {rel}")
        path = Path(
            hf_hub_download(
                repo_id=BASE_DATASET_REPO,
                filename=rel,
                repo_type="dataset",
                token=token,
                local_dir=str(cache_dir),
            )
        )
        # hf bazen local_dir/rel yazar
        candidate = cache_dir / rel
        if candidate.exists():
            path = candidate
        out[rel] = path
        log_fn(f"[faz2] hazir: {path} ({path.stat().st_size // (1024 * 1024)} MB)")
    meta = cache_dir / "faz2_weight_manifest.json"
    meta.write_text(
        json.dumps({k: str(v) for k, v in out.items()}, indent=2),
        encoding="utf-8",
    )
    return out


def _load_safetensors(path: Path) -> dict[str, torch.Tensor]:
    try:
        from safetensors.torch import load_file

        return load_file(str(path))
    except Exception:
        # fallback: torch.load for .ckpt/.pt
        return torch.load(path, map_location="cpu", weights_only=True)


def pick_lora_targets(
    state: dict[str, torch.Tensor],
    *,
    max_matrices: int = 8,
    min_in: int = 64,
    max_in: int = 8192,
) -> list[tuple[str, torch.Tensor]]:
    """2D weight matrislerinden LoRA adaylarini sec (buyukten kucuge)."""
    cands: list[tuple[str, torch.Tensor, int]] = []
    for key, tensor in state.items():
        if not key.endswith("weight"):
            continue
        if tensor.ndim != 2:
            continue
        out_f, in_f = int(tensor.shape[0]), int(tensor.shape[1])
        if in_f < min_in or in_f > max_in:
            continue
        if out_f < 16:
            continue
        cands.append((key, tensor.detach().float().cpu(), out_f * in_f))
    cands.sort(key=lambda x: x[2], reverse=True)
    return [(k, t) for k, t, _ in cands[:max_matrices]]


class FrozenLinearLoRA(nn.Module):
    """W frozen + LoRA(A,B). y = x @ W.T + scale * (x @ A.T) @ B.T"""

    def __init__(
        self,
        weight: torch.Tensor,
        *,
        rank: int = 8,
        scale: float = 1.0,
        name: str = "",
    ) -> None:
        super().__init__()
        out_f, in_f = weight.shape
        self.name = name
        self.in_features = in_f
        self.out_features = out_f
        self.scale = scale
        self.register_buffer("weight", weight.contiguous())
        self.lora_A = nn.Parameter(torch.zeros(rank, in_f))
        self.lora_B = nn.Parameter(torch.zeros(out_f, rank))
        nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
        nn.init.zeros_(self.lora_B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base = nn.functional.linear(x, self.weight)
        delta = (x @ self.lora_A.T) @ self.lora_B.T
        return base + self.scale * delta


class BaseLoRATower(nn.Module):
    """Birden fazla frozen+LoRA katmani; giris projesi ile dim hizalama."""

    def __init__(
        self,
        targets: list[tuple[str, torch.Tensor]],
        *,
        input_dim: int,
        rank: int = 8,
        out_dim: int = 512,
    ) -> None:
        super().__init__()
        if not targets:
            raise ValueError("LoRA hedefi yok")
        layers = nn.ModuleList()
        # Ilk katman in_features'a proj_in
        first_in = int(targets[0][1].shape[1])
        self.proj_in = nn.Linear(input_dim, first_in)
        prev_out = first_in
        for name, w in targets:
            w = w.contiguous()
            # Zincir: onceki out != bu in ise ara projeksiyon
            if prev_out != int(w.shape[1]):
                layers.append(nn.Linear(prev_out, int(w.shape[1])))
                layers.append(nn.GELU())
            layers.append(FrozenLinearLoRA(w, rank=rank, name=name))
            layers.append(nn.GELU())
            prev_out = int(w.shape[0])
        self.layers = layers
        self.proj_out = nn.Linear(prev_out, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.proj_in(x)
        for layer in self.layers:
            h = layer(h)
        return self.proj_out(h)

    def lora_parameters(self) -> list[nn.Parameter]:
        params: list[nn.Parameter] = []
        for m in self.modules():
            if isinstance(m, FrozenLinearLoRA):
                params.extend([m.lora_A, m.lora_B])
        return params