from __future__ import annotations

import gc
import hashlib
import importlib.metadata
import json
import logging
import os
from pathlib import Path

import mlx.core as mx
from transformers import AutoTokenizer

from .ops import attention, rms_norm, rotate, silu
from .weights import Linear, WeightStore

log = logging.getLogger(__name__)


class TextEncoder:
    """Qwen3-VL's first fifty language layers, with no final RMS normalization.

    Layers stream one at a time. Conditioning is cached by weights, exact token
    presentation and precision policy; it is independent of diffusion seed.
    """

    def __init__(
        self, path, *, cache_dir="cache", bits=None, group_size=64, query_chunk=0
    ):
        self.path = Path(path)
        self.config = json.loads((self.path / "config.json").read_text())
        self.text_config = self.config["text_config"]
        self.store = WeightStore(self.path)
        self.tokenizer = AutoTokenizer.from_pretrained(
            str(self.path), local_files_only=True
        )
        self.cache_dir = Path(cache_dir)
        self.bits, self.group_size, self.query_chunk = bits, group_size, query_chunk

    def encode(self, prompt, *, images=(), on_layer=None):
        self.cache_hit = False
        if not prompt.strip():
            raise ValueError("Prompt must be nonempty")
        if images:
            from .vision import image_presentation

            ids, tags, positions, vision = image_presentation(self, prompt, images)
        else:
            raw = self.tokenizer.encode(prompt, add_special_tokens=False)[:4096]
            ids = raw + [0] * (-len(raw) % 8)
            tags = mx.ones(len(raw), mx.int32)
            positions = None
            vision = None
        identity = {
            "version": 5,
            "weights": self.store.fingerprint(),
            "config": self.config,
            "versions": {
                p: importlib.metadata.version(p)
                for p in ("mlx", "numpy", "Pillow", "transformers")
            },
            "ids": ids,
            "bits": self.bits,
            "group_size": self.group_size,
            "query_chunk": self.query_chunk,
            "tap": 50,
            "tf32": os.environ.get("MLX_ENABLE_TF32"),
            "images": [
                hashlib.sha256(im.tobytes()).hexdigest() + str(im.size) for im in images
            ],
        }
        if images:
            identity["vision_layer_norm_revision"] = 2
        key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest()[
            :24
        ]
        target = self.cache_dir / f"conditioning-{key}.safetensors"
        if target.exists():
            data = mx.load(str(target))
            mx.eval(data)
            self.cache_hit = True
            log.info(
                "Conditioning cache hit: %d effective tokens", data["hidden"].shape[0]
            )
            return data["hidden"], data["tags"]
        prefix = "model.language_model"
        x = self.store.get(prefix + ".embed_tokens.weight", rows=ids)[None]
        mx.eval(x)
        if vision is not None:
            from .vision import encode_vision

            features, deep = encode_vision(
                self.store, self.config["vision_config"], vision
            )
            image_indices = mx.array(
                [i for i, t in enumerate(ids) if t == self.config["image_token_id"]]
            )
            x[0, image_indices] = features
        cfg = self.text_config
        dim, nq, nk = (
            cfg["head_dim"],
            cfg["num_attention_heads"],
            cfg["num_key_value_heads"],
        )
        if positions is None:
            positions = mx.broadcast_to(
                mx.arange(len(ids), dtype=mx.float32)[None], (3, len(ids))
            )
        else:
            positions = mx.array(positions, mx.float32)
        inv = 1 / (cfg["rope_theta"] ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
        angles = positions[:, :, None] * inv[None, None]
        # Qwen3-VL interleaves spatial axes into the temporal frequency stream.
        sections = cfg["rope_scaling"]["mrope_section"]
        freq = angles[0]
        for axis, offset in ((1, 1), (2, 2)):
            stop = sections[axis] * 3
            freq[:, offset:stop:3] = angles[axis, :, offset:stop:3]
        freq = mx.concatenate((freq, freq), axis=-1)
        rope = (
            mx.cos(freq).astype(mx.bfloat16)[None, :, None],
            mx.sin(freq).astype(mx.bfloat16)[None, :, None],
        )
        mx.eval(rope)
        for i in range(50):
            layer = f"{prefix}.layers.{i}"

            def linear(suffix, layer=layer):
                return Linear.load(
                    self.store,
                    layer + "." + suffix,
                    bits=self.bits,
                    group_size=self.group_size,
                )

            qproj, kproj, vproj = (
                linear("self_attn." + p + "_proj") for p in ("q", "k", "v")
            )
            oproj = linear("self_attn.o_proj")
            h = rms_norm(
                x,
                self.store.get(layer + ".input_layernorm.weight"),
                cfg["rms_norm_eps"],
                rounded=True,
            )
            q = qproj(h).reshape(1, len(ids), nq, dim)
            k = kproj(h).reshape(1, len(ids), nk, dim)
            v = vproj(h).reshape(1, len(ids), nk, dim)
            q = rms_norm(
                q,
                self.store.get(layer + ".self_attn.q_norm.weight"),
                cfg["rms_norm_eps"],
                rounded=True,
            )
            k = rms_norm(
                k,
                self.store.get(layer + ".self_attn.k_norm.weight"),
                cfg["rms_norm_eps"],
                rounded=True,
            )
            q, k = rotate(q, *rope), rotate(k, *rope)
            h = attention(q, k, v, causal=True, query_chunk=self.query_chunk).reshape(
                1, len(ids), -1
            )
            x = x + oproj(h)
            mx.eval(x)
            del qproj, kproj, vproj, oproj, q, k, v, h
            gate, up, down = (
                linear("mlp." + p + "_proj") for p in ("gate", "up", "down")
            )
            h = rms_norm(
                x,
                self.store.get(layer + ".post_attention_layernorm.weight"),
                cfg["rms_norm_eps"],
                rounded=True,
            )
            x = x + down(silu(gate(h)) * up(h))
            if vision is not None and i < len(deep):
                x = x.at[0, image_indices].add(deep[i])
            mx.eval(x)
            del gate, up, down, h
            if on_layer:
                on_layer(i, x)
            if (i + 1) % 10 == 0:
                log.info("Text encoder %d/50", i + 1)
            mx.clear_cache()
        hidden = x[0, : len(tags)].astype(mx.bfloat16)
        mx.eval(hidden)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        tmp = target.with_suffix(".tmp.safetensors")
        mx.save_safetensors(
            str(tmp),
            {"hidden": hidden, "tags": tags},
            metadata={"identity": json.dumps(identity)},
        )
        tmp.replace(target)
        gc.collect()
        mx.clear_cache()
        return hidden, tags