"""Latent adapter R&D scaffold (requires real TRELLIS/Hunyuan model tensors)."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

try:
    import torch
    import torch.nn as nn
except ImportError:  # pragma: no cover - optional for preprocess-only runs
    torch = None  # type: ignore
    nn = None  # type: ignore


def _utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


@dataclass
class LatentShapeLogger:
    """Log tensor shapes until real model hooks are wired."""

    log_path: Path | None = None
    records: list[dict[str, Any]] = field(default_factory=list)

    def log(self, name: str, tensor: Any, *, source: str = "unknown") -> dict[str, Any]:
        if torch is not None and isinstance(tensor, torch.Tensor):
            shape = list(tensor.shape)
            dtype = str(tensor.dtype)
            device = str(tensor.device)
        elif hasattr(tensor, "shape"):
            shape = list(getattr(tensor, "shape"))
            dtype = str(getattr(tensor, "dtype", "?"))
            device = "numpy"
        else:
            shape = []
            dtype = type(tensor).__name__
            device = "?"

        entry = {
            "name": name,
            "source": source,
            "shape": shape,
            "dtype": dtype,
            "device": device,
            "logged_at": _utc_now(),
        }
        self.records.append(entry)
        if self.log_path is not None:
            self.log_path.parent.mkdir(parents=True, exist_ok=True)
            with open(self.log_path, "a", encoding="utf-8") as handle:
                handle.write(json.dumps(entry, ensure_ascii=False) + "\n")
        return entry

    def summary(self) -> dict[str, Any]:
        return {
            "count": len(self.records),
            "shapes": {r["name"]: r["shape"] for r in self.records},
            "updated_at": _utc_now(),
        }


if nn is not None:

    class TrellisHunyuanLatentAdapter(nn.Module):
        """Small MLP adapter: TRELLIS structured latent -> Hunyuan conditioning latent."""

        def __init__(
            self,
            trellis_dim: int,
            hunyuan_dim: int,
            *,
            hidden_dim: int = 512,
            depth: int = 2,
            adapter_type: str = "mlp",
        ) -> None:
            super().__init__()
            self.adapter_type = adapter_type
            self.trellis_dim = trellis_dim
            self.hunyuan_dim = hunyuan_dim
            layers: list[nn.Module] = []
            in_dim = trellis_dim
            for _ in range(max(depth - 1, 0)):
                layers.extend([nn.Linear(in_dim, hidden_dim), nn.GELU()])
                in_dim = hidden_dim
            layers.append(nn.Linear(in_dim, hunyuan_dim))
            self.net = nn.Sequential(*layers)
            self.logger = LatentShapeLogger()

        def forward(self, trellis_latent: "torch.Tensor") -> "torch.Tensor":
            self.logger.log("trellis_latent_in", trellis_latent, source="trellis")
            flat = trellis_latent.reshape(trellis_latent.shape[0], -1)
            if flat.shape[-1] != self.trellis_dim:
                raise ValueError(
                    f"Expected trellis_dim={self.trellis_dim}, got {flat.shape[-1]}"
                )
            out = self.net(flat)
            self.logger.log("hunyuan_latent_out", out, source="adapter")
            return out

else:

    class TrellisHunyuanLatentAdapter:  # type: ignore[no-redef]
        def __init__(self, *args, **kwargs) -> None:
            raise ImportError("torch required for TrellisHunyuanLatentAdapter")


def research_notes() -> dict[str, Any]:
    """Document planned losses and integration points for future real-model wiring."""
    return {
        "status": "rd_only",
        "requires": [
            "trellis_structured_latent_export",
            "hunyuan_denoising_conditioning_hook",
        ],
        "candidate_losses": [
            "texture_view_reconstruction",
            "uv_consistency",
            "normal_curvature_texture_consistency",
        ],
        "adapter_variants": ["linear", "mlp", "cross_attention"],
    }