#!/usr/bin/env python3
import importlib.util
import json
import sys
from pathlib import Path
from types import SimpleNamespace

import torch
from exllamav3 import Config, Model, Tokenizer
from safetensors.torch import load_file


def load_joint_module(path):
    spec = importlib.util.spec_from_file_location("clef_joint_schema_model", path / "joint_schema_model.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


class TokenizerAdapter:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def __call__(self, text, add_special_tokens=False):
        ids = self.tokenizer.encode(text, encode_special_tokens=True)
        return SimpleNamespace(input_ids=ids[0].tolist())


class ClefEXL3:
    def __init__(self, model_path, device="cuda:0"):
        self.path = Path(model_path)
        self.device = torch.device(device)
        self.joint = load_joint_module(self.path)
        config = Config.from_directory(str(self.path))
        self.tokenizer = TokenizerAdapter(Tokenizer.from_config(config))
        self.model = Model.from_config(config, component="text")
        self.model.load(progressbar=True, device=device)
        head_config = json.loads((self.path / "joint_head_config.json").read_text())
        self.head = self.joint.JointSchemaHead(**head_config)
        self.head.load_state_dict(load_file(self.path / "joint_head.safetensors"), strict=True)
        self.head = self.head.to(device=self.device, dtype=torch.float16).eval()
        lm_head = self.model.modules[self.model.logit_layer_idx]
        if getattr(lm_head, "quant_type", None) != "fp16":
            raise RuntimeError("Clef EXL3 requires an FP16 LM head; quantize with --head_bits 16")
        raw = json.loads((self.path / "config.json").read_text())
        text_config = raw.get("text_config", raw)
        vocab_size = int(text_config["vocab_size"])
        self.output_embedding_weight = lm_head.inner.weight.T[:vocab_size]

    @torch.inference_mode()
    def hidden_states(self, input_ids):
        params = {"attn_mode": "flash_attn_nc"}
        x = self.model.prepare_inputs(input_ids, params)
        for module, instance, _ in self.model.fwd_modules:
            if module.caps.get("logits_output"):
                break
            params["layer_instance"] = instance
            x = module.prepare_for_device(x, params)
            x = module.forward(x, params)
        return x

    @torch.inference_mode()
    def systemone(self, request, max_length=16384):
        if request.get("images") or request.get("videos"):
            raise NotImplementedError("Clef EXL3 adapter currently supports text/JSON state only")
        encoded = self.joint.encode_record(
            self.tokenizer,
            request,
            max_length=max_length,
            processor=None,
        )
        input_ids = torch.tensor([encoded.input_ids], dtype=torch.long, device=self.device)
        attention_mask = torch.ones_like(input_ids)
        hidden = self.hidden_states(input_ids).to(torch.float16)
        logits = self.head(
            hidden,
            input_ids,
            attention_mask,
            [encoded],
            self.output_embedding_weight,
        )[0]
        questions = request["questions"]
        answers = {}
        for question, question_logits in zip(encoded.questions, logits):
            probabilities = dict(zip(question.option_ids, question_logits.float().softmax(-1).tolist()))
            answers[question.question_id] = self.joint.systemone_answer(
                questions[question.question_id], probabilities
            )
        return {
            "model": request.get("model", "clef-flash-exl3"),
            "answers": answers,
            "usage": {"input_tokens": len(encoded.input_ids), "output_tokens": 0},
        }