import os
import gc
import threading
from typing import Optional

import torch
from llama_cpp import Llama
from comfy import model_management

import folder_paths

# -------------------------------------------------------------------
# 1. Helpers for GGUF file listing & resolution (unchanged)
# -------------------------------------------------------------------
_CACHE = {}
_CACHE_LOCK = threading.Lock()

def _list_gguf_files():
    candidates = []
    base_dirs = []
    try:
        base_dirs.extend(folder_paths.get_folder_paths("text_encoders"))
    except Exception:
        pass

    for base in base_dirs:
        if not base or not os.path.isdir(base):
            continue
        for root, _, files in os.walk(base):
            for name in files:
                if name.lower().endswith(".gguf"):
                    full = os.path.join(root, name)
                    try:
                        rel = os.path.relpath(full, base)
                        if rel not in candidates:
                            candidates.append(rel)
                    except Exception:
                        if full not in candidates:
                            candidates.append(full)

    if "<manual_path>" not in candidates:
        candidates.append("<manual_path>")
    return candidates

def _resolve_model_path(model_name: str, manual_model_path: str = "") -> str:
    if os.path.isabs(model_name) and os.path.isfile(model_name):
        return model_name
    if model_name == "<manual_path>":
        p = (manual_model_path or "").strip()
        if not p:
            raise ValueError("model is <manual_path> but manual_model_path is empty")
        if not os.path.isfile(p):
            raise FileNotFoundError(f"GGUF model not found: {p}")
        return p
    full = folder_paths.get_full_path("text_encoders", model_name)
    if full and os.path.isfile(full):
        return full
    try:
        for base in folder_paths.get_folder_paths("text_encoders"):
            probe = os.path.join(base, model_name)
            if os.path.isfile(probe):
                return probe
    except Exception:
        pass
    raise FileNotFoundError(f"Could not resolve GGUF model path for: {model_name}")

def _build_messages(system_prompt: str, user_prompt: str):
    messages = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    messages.append({"role": "user", "content": user_prompt.strip()})
    return messages

def _assemble_prompt(messages: list, model_family: str) -> str:
    prompt = ""
    if model_family == "granite":
        for msg in messages:
            prompt += f"<|start_of_role|>{msg['role']}<|end_of_role|>{msg['content']}<|end_of_text|>\n"
        prompt += "<|start_of_role|>assistant<|end_of_role|><think>\n"
    else:
        for msg in messages:
            prompt += f"<|im_start|>{msg['role']}\n{msg['content']}<|im_end|>\n"
        prompt += "<|im_start|>assistant\n<think>"
    return prompt


# -------------------------------------------------------------------
# 2. Lookup table for Mordant models (unchanged)
# -------------------------------------------------------------------
MORDANT_INFO = {
    "mordant-1.2b": {
        "layers": 16,
        "bf16_total_gb": 2.35,
        "per_layer": {
            "bf16": 2.35 / 16,
            "q8_0": 1.27 / 16,
            "q6_k_m": 1.01682 / 16,
            "q5_k_m": None,
            "q4_k_m": None,
        }
    },
    "mordant-3b": {
        "layers": 40,
        "bf16_total_gb": 6.99,
        "per_layer": {
            "bf16": 6.99 / 40,
            "q8_0": 3.86 / 40,
            "q6_k_m": 3.05 / 40,
            "q5_k_m": 2.70 / 40,
            "q4_k_m": None,
        }
    },
    "mordant-7b": {
        "layers": 32,
        "bf16_total_gb": 16.29,
        "per_layer": {
            "bf16": 16.29 / 32,
            "q8_0": 9.58 / 32,
            "q6_k_m": 7.85 / 32,
            "q5_k_m": 7.08 / 32,
            "q4_k_m": None,
        }
    },
    "mordant-12b": {
        "layers": 40,
        "bf16_total_gb": 24.81,
        "per_layer": {
            "bf16": 24.81 / 40,
            "q8_0": 13.55 / 40,
            "q6_k_m": 10.64 / 40,
            "q5_k_m": 9.33 / 40,
            "q4_k_m": 8.11 / 40,
        }
    },
}

def _detect_model_info(model_path: str):
    try:
        llm_meta = Llama(model_path=model_path, vocab_only=True, n_gpu_layers=0, verbose=False)
        meta = llm_meta.metadata
        arch = meta.get("general.architecture", "unknown")
        total_layers = int(meta.get(f"{arch}.block_count", 0))
        quant_version = meta.get("general.quantization_version", "unknown").lower()
        if quant_version == "unknown":
            quant_version = meta.get("tokenizer.ggml.quantization_version", "unknown").lower()
        del llm_meta
    except Exception:
        fname = os.path.basename(model_path).lower()
        for name in MORDANT_INFO:
            if name in fname:
                info = MORDANT_INFO[name]
                total_layers = info["layers"]
                quant_version = "bf16"
                for q in ["q8_0", "q6_k_m", "q5_k_m", "q4_k_m", "q3_k_m", "q2_k_m"]:
                    if q in fname:
                        quant_version = q
                        break
                break
        else:
            total_layers = 32
            quant_version = "bf16"

    model_name = None
    fname_lower = os.path.basename(model_path).lower()
    for name in MORDANT_INFO:
        if name in fname_lower:
            model_name = name
            break

    if model_name is None:
        file_size_gb = os.path.getsize(model_path) / (1024**3)
        per_layer_vram = file_size_gb / total_layers if total_layers > 0 else 0.5
        return (model_name or "unknown", total_layers, per_layer_vram, quant_version)

    info = MORDANT_INFO[model_name]
    layers = info["layers"]
    per_layer_dict = info.get("per_layer", {})
    per_layer_vram = None
    quant_clean = quant_version.replace("_0", "").replace("_1", "").lower()
    if quant_version in per_layer_dict and per_layer_dict[quant_version] is not None:
        per_layer_vram = per_layer_dict[quant_version]
    elif quant_clean in per_layer_dict and per_layer_dict[quant_clean] is not None:
        per_layer_vram = per_layer_dict[quant_clean]
    else:
        bf16_per_layer = per_layer_dict.get("bf16")
        if bf16_per_layer is None:
            bf16_total = info["bf16_total_gb"]
            bf16_per_layer = bf16_total / layers
        actual_file_size_gb = os.path.getsize(model_path) / (1024**3)
        per_layer_vram = actual_file_size_gb / layers

    return model_name, layers, per_layer_vram, quant_version


def _auto_n_gpu_layers(model_path: str, verbose: bool = False) -> int:
    if not torch.cuda.is_available():
        return 0

    total_vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    usable_vram_gb = max(0, total_vram_gb - 2.0)

    model_name, max_layers, per_layer_vram, quant = _detect_model_info(model_path)
    if per_layer_vram <= 0:
        return 0

    n_layers = int(usable_vram_gb / per_layer_vram)
    n_layers = max(0, min(max_layers, n_layers))

    if verbose:
        print(f"[Mordant Enhancer] GPU VRAM: {total_vram_gb:.2f} GB, "
              f"Reserved: 2.00 GB, Usable: {usable_vram_gb:.2f} GB")
        print(f"[Mordant Enhancer] Model: {model_name}, Layers: {max_layers}, "
              f"Quant: {quant}, Per‑layer VRAM: {per_layer_vram:.4f} GB")
        print(f"[Mordant Enhancer] Auto‑selected GPU layers: {n_layers}")

    return n_layers


# -------------------------------------------------------------------
# 3. Main node class – Mordant Prompt Enhancer
# -------------------------------------------------------------------
class MordantPromptEnhancer:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": (_list_gguf_files(),),
                "text": ("STRING", {"multiline": True, "default": "Rewrite this image prompt conservatively."}),
                "mode": (["enhancer", "general"], {"default": "enhancer"}),
                "enable_enhancement": ("BOOLEAN", {"default": True}),
                "sampling_mode": ("BOOLEAN", {"default": True}),
                "max_new_tokens": ("INT", {"default": 8192, "min": 1, "max": 8192}),
                "seed": ("INT", {"default": 0, "min": -1, "max": 2147483647}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 5.0, "step": 0.05}),
                "top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
                "top_k": ("INT", {"default": 50, "min": 0, "max": 1000}),
                "repeat_penalty": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.01}),
                "n_batch": ("INT", {"default": 2048, "min": 32, "max": 4096, "step": 32}),
                "keep_loaded": ("BOOLEAN", {"default": False}),
                "offload_kqv": ("BOOLEAN", {"default": True}),
                "flash_attn": ("BOOLEAN", {"default": True}),
                "verbose": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "manual_model_path": ("STRING", {"default": "", "multiline": False}),
                "custom_system_prompt": ("STRING", {"default": "", "multiline": True}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("composition", "thinking")
    FUNCTION = "generate"
    CATEGORY = "text/llm"

    def generate(self, model, text, mode, enable_enhancement,
                 sampling_mode, max_new_tokens, seed, temperature, top_p, top_k,
                 repeat_penalty, n_batch, keep_loaded, offload_kqv, flash_attn,
                 verbose, manual_model_path="", custom_system_prompt=""):

        if not enable_enhancement:
            return (text, "")

        model_path = _resolve_model_path(model, manual_model_path)
        system_prompt = (custom_system_prompt or "").strip()

        n_ctx = 8192
        n_threads = os.cpu_count() or 4
        n_gpu_layers = _auto_n_gpu_layers(model_path, verbose=verbose)

        if not sampling_mode:
            eff_temp, eff_top_p, eff_top_k, eff_repeat_penalty = 0.2, 0.9, 20, max(repeat_penalty, 1.05)
        else:
            eff_temp, eff_top_p, eff_top_k, eff_repeat_penalty = temperature, top_p, top_k, repeat_penalty

        try:
            from llama_cpp import Llama
        except ImportError:
            raise RuntimeError("llama-cpp-python not installed.")

        def _create_model():
            kwargs = dict(
                model_path=model_path,
                n_ctx=n_ctx,
                n_gpu_layers=n_gpu_layers,
                n_batch=n_batch,
                verbose=verbose,
                cuda_graphs=False,
                use_mmap=False,
                use_mlock=False,
            )
            if n_threads is not None:
                kwargs["n_threads"] = n_threads
            try:
                kwargs["offload_kqv"] = offload_kqv
                kwargs["flash_attn"] = flash_attn
                return Llama(**kwargs)
            except TypeError:
                kwargs.pop("offload_kqv", None)
                kwargs.pop("flash_attn", None)
                return Llama(**kwargs)

        if keep_loaded:
            key = (model_path, n_ctx, n_gpu_layers, n_threads, n_batch, offload_kqv, flash_attn)
            with _CACHE_LOCK:
                llm = _CACHE.get(key)
                if llm is None:
                    llm = _create_model()
                    _CACHE[key] = llm
        else:
            llm = _create_model()

        actual_family = "chatml"
        if hasattr(llm, 'metadata'):
            chat_template = llm.metadata.get('tokenizer.chat_template', '').lower()
            if 'granite' in chat_template or 'start_of_role' in chat_template:
                actual_family = "granite"

        messages = _build_messages(system_prompt, text)
        prompt = _assemble_prompt(messages, actual_family)

        out = llm.create_completion(
            prompt=prompt,
            max_tokens=max_new_tokens,
            temperature=eff_temp,
            top_p=eff_top_p,
            top_k=eff_top_k,
            repeat_penalty=eff_repeat_penalty,
            seed=None if seed < 0 else seed,
        )

        result_text = out["choices"][0]["text"]
        if "</think>" in (result_text or ""):
            parts = (result_text or "").rsplit("</think>", 1)
            analysis_text, final_text = parts[0].strip(), parts[1].strip()
        else:
            analysis_text, final_text = "", (result_text or "").strip()

        return (final_text, analysis_text)


# -------------------------------------------------------------------
# 4. Node registration
# -------------------------------------------------------------------
NODE_CLASS_MAPPINGS = {
    "MordantPromptEnhancer": MordantPromptEnhancer,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MordantPromptEnhancer": "Mordant Prompt Enhancer",
}