# coding=utf-8
"""Lazy Qwen3.5 visual-block sidecar for Swarm-MoE runtime vision routes.

The sidecar is intentionally outside ``forward`` and ``generate``. A normal HF text
load keeps the original Swarm-MoE state dict unchanged; browser/vision tooling may
opt in and load these transferred visual tensors only when it needs pixel grounding.
"""

from __future__ import annotations

import json
import math
import re
import shutil
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F


VISUAL_PREFIX = "prime_adapter.model.model.visual."
DUAL_RETINA_KEY = "dual_retina.shared_retina.patch.weight"
VISION_TOKEN_NAMES = (
    "<|object_ref_start|>",
    "<|object_ref_end|>",
    "<|box_start|>",
    "<|box_end|>",
    "<|quad_start|>",
    "<|quad_end|>",
    "<|vision_start|>",
    "<|vision_end|>",
    "<|vision_pad|>",
    "<|image_pad|>",
    "<|video_pad|>",
)


def _layer_norm(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, eps: float = 1e-6) -> torch.Tensor:
    y = (x.float() - x.float().mean(dim=-1, keepdim=True))
    y = y * torch.rsqrt(y.pow(2).mean(dim=-1, keepdim=True) + eps)
    y = y.to(dtype=x.dtype) * weight
    return y if bias is None else y + bias


def _linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
    return F.linear(x, weight, bias)


def _token_map_from_tokenizer(tokenizer_json: Path, target_vocab_size: int) -> dict[str, Any]:
    if not tokenizer_json.exists():
        return {"status": "missing_source_tokenizer", "tokens": {}}
    data = json.loads(tokenizer_json.read_text(encoding="utf-8"))
    added = data.get("added_tokens") or []
    by_content = {str(row.get("content")): row for row in added if isinstance(row, dict)}
    tokens = {}
    for content in VISION_TOKEN_NAMES:
        row = by_content.get(content)
        if row:
            tokens[content] = {
                "source_id": row.get("id"),
                "target_id": None,
                "encoding_policy": "external_processor_marker_not_embedded",
            }
    return {
        "status": "metadata_only_no_embedding_resize",
        "target_vocab_size": int(target_vocab_size),
        "reason": "source vision token ids are outside the 32k tied embedding table; pixel tensors route through the sidecar instead",
        "tokens": tokens,
    }


def export_qwen35_vision_sidecar(
    source: str | Path,
    destination: str | Path,
    *,
    target_vocab_size: int = 32000,
    include_source_path: bool = False,
) -> dict[str, Any]:
    """Copy Qwen3.5 visual block tensors into a compact Swarm-MoE sidecar folder."""

    from safetensors import safe_open
    from safetensors.torch import save_file

    source = Path(source)
    destination = Path(destination)
    destination.mkdir(parents=True, exist_ok=True)
    index_path = source / "model.safetensors.index.json"
    if not index_path.exists():
        raise FileNotFoundError(f"missing safetensors index: {index_path}")

    index = json.loads(index_path.read_text(encoding="utf-8"))
    weight_map: dict[str, str] = index.get("weight_map") or {}
    keys = [k for k in weight_map if k.startswith(VISUAL_PREFIX)]
    if DUAL_RETINA_KEY in weight_map:
        keys.append(DUAL_RETINA_KEY)
    if not keys:
        raise ValueError("no Qwen3.5 visual keys found in source checkpoint")

    tensors: dict[str, torch.Tensor] = {}
    by_shard: dict[str, list[str]] = {}
    for key in keys:
        by_shard.setdefault(weight_map[key], []).append(key)
    for shard, shard_keys in sorted(by_shard.items()):
        with safe_open(source / shard, framework="pt", device="cpu") as handle:
            for key in shard_keys:
                tensors[key] = handle.get_tensor(key).contiguous()

    source_label = str(source) if include_source_path else "external_source_redacted_for_portable_hf_package"
    save_file(tensors, destination / "model.safetensors", metadata={"format": "pt", "source": source_label})

    copied_files = []
    for name in ("preprocessor_config.json", "video_preprocessor_config.json"):
        src_file = source / name
        if src_file.exists():
            shutil.copy2(src_file, destination / name)
            copied_files.append(name)

    src_config = json.loads((source / "config.json").read_text(encoding="utf-8"))
    token_map = _token_map_from_tokenizer(source / "tokenizer.json", target_vocab_size)
    (destination / "vision_token_map.json").write_text(json.dumps(token_map, indent=2), encoding="utf-8")

    block_ids = sorted({int(m.group(1)) for k in keys for m in [re.search(r"\.blocks\.(\d+)\.", k)] if m})
    manifest = {
        "format": "phill_swarm_qwen35_vision_sidecar_v1",
        "source_checkpoint": source_label,
        "tensor_file": "model.safetensors",
        "tensor_count": len(tensors),
        "visual_tensor_count": sum(1 for k in tensors if k.startswith(VISUAL_PREFIX)),
        "dual_retina_tensor_count": 1 if DUAL_RETINA_KEY in tensors else 0,
        "block_count": len(block_ids),
        "block_ids": block_ids,
        "dtype": "bfloat16",
        "prime_vision_config": src_config.get("prime_vision_config") or {},
        "copied_processor_files": copied_files,
        "token_map_file": "vision_token_map.json",
        "hf_compatibility": "runtime sidecar only; SwarmMoEForCausalLM.forward/generate signatures are unchanged",
        "memory_policy": "not loaded until route_vla_action is called with vision_sidecar_enabled=True",
    }
    (destination / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    return manifest


class Qwen35VisionSidecar:
    """Small runtime executor for transferred Qwen3.5 visual tensors."""

    def __init__(self, sidecar_dir: str | Path, *, device: torch.device | str = "cpu", dtype: torch.dtype | None = None, max_blocks: int = 24):
        self.sidecar_dir = Path(sidecar_dir)
        self.device = torch.device(device)
        self.dtype = dtype or torch.bfloat16
        self.max_blocks = max(0, int(max_blocks))
        self.manifest = json.loads((self.sidecar_dir / "manifest.json").read_text(encoding="utf-8"))
        self._tensors: dict[str, torch.Tensor] | None = None

    @property
    def available(self) -> bool:
        return (self.sidecar_dir / "model.safetensors").exists()

    def _load(self) -> dict[str, torch.Tensor]:
        if self._tensors is not None:
            return self._tensors
        from safetensors import safe_open

        tensors: dict[str, torch.Tensor] = {}
        with safe_open(self.sidecar_dir / "model.safetensors", framework="pt", device="cpu") as handle:
            for key in handle.keys():
                tensors[key] = handle.get_tensor(key).to(device=self.device, dtype=self.dtype)
        self._tensors = tensors
        return tensors

    def _get(self, suffix: str, *, required: bool = True) -> torch.Tensor | None:
        key = suffix if suffix.startswith(VISUAL_PREFIX) or suffix == DUAL_RETINA_KEY else VISUAL_PREFIX + suffix
        tensors = self._load()
        value = tensors.get(key)
        if value is None and required:
            raise KeyError(key)
        return value

    def _block(self, x: torch.Tensor, idx: int, heads: int) -> torch.Tensor:
        prefix = f"blocks.{idx}."
        n1w = self._get(prefix + "norm1.weight")
        n1b = self._get(prefix + "norm1.bias", required=False)
        h = _layer_norm(x, n1w, n1b)
        qkv = _linear(h, self._get(prefix + "attn.qkv.weight"), self._get(prefix + "attn.qkv.bias", required=False))
        bsz, seq, dim3 = qkv.shape
        head_dim = dim3 // (3 * heads)
        qkv = qkv.view(bsz, seq, 3, heads, head_dim).permute(2, 0, 3, 1, 4)
        attn = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2], dropout_p=0.0, is_causal=False)
        attn = attn.transpose(1, 2).reshape(bsz, seq, heads * head_dim)
        x = x + _linear(attn, self._get(prefix + "attn.proj.weight"), self._get(prefix + "attn.proj.bias", required=False))

        n2w = self._get(prefix + "norm2.weight")
        n2b = self._get(prefix + "norm2.bias", required=False)
        h = _layer_norm(x, n2w, n2b)
        h = F.gelu(_linear(h, self._get(prefix + "mlp.linear_fc1.weight"), self._get(prefix + "mlp.linear_fc1.bias", required=False)), approximate="tanh")
        return x + _linear(h, self._get(prefix + "mlp.linear_fc2.weight"), self._get(prefix + "mlp.linear_fc2.bias", required=False))

    @torch.no_grad()
    def encode(self, image_tensor: torch.Tensor, *, max_blocks: int | None = None) -> dict[str, Any]:
        if image_tensor.ndim == 3:
            image_tensor = image_tensor.unsqueeze(0)
        if image_tensor.shape[1] not in {1, 3, 4} and image_tensor.shape[-1] in {1, 3, 4}:
            image_tensor = image_tensor.permute(0, 3, 1, 2).contiguous()
        x = image_tensor[:, :3].to(device=self.device, dtype=self.dtype)
        if x.max() > 2:
            x = x / 255.0
        x = F.interpolate(x.float(), size=(224, 224), mode="bilinear", align_corners=False).to(dtype=self.dtype)
        x = (x - 0.5) / 0.5
        x = x.unsqueeze(2).repeat(1, 1, 2, 1, 1)
        x = F.conv3d(x, self._get("patch_embed.proj.weight"), self._get("patch_embed.proj.bias", required=False), stride=(2, 16, 16))
        x = x.flatten(2).transpose(1, 2).contiguous()

        pos = self._get("pos_embed.weight", required=False)
        if pos is not None and pos.shape[0] >= x.shape[1]:
            x = x + pos[: x.shape[1]].unsqueeze(0).to(dtype=x.dtype)

        cfg = self.manifest.get("prime_vision_config") or {}
        heads = int(cfg.get("num_heads") or 16)
        blocks = min(int(max_blocks if max_blocks is not None else self.max_blocks), int(self.manifest.get("block_count") or 0))
        for idx in range(blocks):
            x = self._block(x, idx, heads)

        pooled = x.float().mean(dim=1)
        return {
            "status": "ok",
            "patches": int(x.shape[1]),
            "hidden_size": int(x.shape[-1]),
            "blocks_used": int(blocks),
            "pooled_norm": float(pooled.norm(dim=-1).mean().item()),
            "pooled_mean": float(pooled.mean().item()),
            "pooled_std": float(pooled.std().item()),
        }

    @torch.no_grad()
    def route(self, image_tensor: torch.Tensor, command: str = "", action_dim: int = 8) -> dict[str, Any]:
        packet = self.encode(image_tensor)
        lower = (command or "").lower()
        visual_energy = min(1.0, float(packet["pooled_std"]) / 2.5)
        wants_click = any(token in lower for token in ("click", "press", "select", "open"))
        wants_type = any(token in lower for token in ("type", "enter", "fill", "search"))
        wants_scroll = any(token in lower for token in ("scroll", "more", "lower", "down"))
        controls = {
            "mouse": {
                "click": float(0.78 if wants_click else 0.18 + visual_energy * 0.12),
                "scroll": float(0.82 if wants_scroll else 0.18 + visual_energy * 0.18),
                "center_x": 0.5,
                "center_y": 0.52,
            },
            "keyboard": {
                "type": float(0.82 if wants_type else 0.05),
                "enter": float(0.45 if wants_type else 0.05),
            },
            "observe": float(0.88 if not (wants_click or wants_type or wants_scroll) else 0.35),
        }
        return {
            "status": "ok",
            "controls": controls,
            "route": {
                "source": "Qwen3.5 visual sidecar",
                "transferred": True,
                "runtime_only": True,
                "weight_file": str(self.sidecar_dir / "model.safetensors"),
                "blocks_used": packet["blocks_used"],
                "patches": packet["patches"],
                "pooled_norm": round(packet["pooled_norm"], 4),
                "pooled_mean": round(packet["pooled_mean"], 4),
                "pooled_std": round(packet["pooled_std"], 4),
                "action_dim": int(action_dim or 0),
            },
        }


__all__ = ["Qwen35VisionSidecar", "export_qwen35_vision_sidecar"]