import multiprocessing
import os
import re
import json
import base64
from io import BytesIO
from PIL import Image
from glob import glob
from rapidfuzz import fuzz, process
import llama_cpp
from llama_cpp import Llama
from comfy_api.latest import io
from pathlib import Path
import torch
import folder_paths
import comfy.model_management

from .utils import *

# Standard GGML enum integer values: 1 = F16, 8 = Q8_0, 2 = Q4_0
KV_CACHE_TYPES = {
    "f16": getattr(getattr(llama_cpp, "llama_cpp", llama_cpp), "GGML_TYPE_F16", 1),
    "q8_0": getattr(getattr(llama_cpp, "llama_cpp", llama_cpp), "GGML_TYPE_Q8_0", 8),
    "q4_0": getattr(getattr(llama_cpp, "llama_cpp", llama_cpp), "GGML_TYPE_Q4_0", 2),
}

class LlamaModel:
    def __init__(self):
        self.llm = None

class LlamaPatcher:
    """Minimal patcher so comfy.model_management can load and unload the llama.cpp model like any other model."""
    def __init__(self, model_file, mmproj_file, n_gpu_layers, llama_kwargs):
        # Stable object tracked by LoadedModel; the Llama instance inside is created and freed on demand
        self.model = LlamaModel()
        self.parent = None
        self.model_file = model_file
        self.mmproj_file = mmproj_file
        self.llama_kwargs = llama_kwargs
        self.size = os.path.getsize(model_file) + (os.path.getsize(mmproj_file) if mmproj_file else 0)
        self.offload_device = torch.device("cpu")
        self.load_device = self.offload_device if n_gpu_layers == 0 else comfy.model_management.get_torch_device()

    def load(self):
        comfy.model_management.load_models_gpu([self])
        return self.model.llm

    def model_size(self):
        return self.size

    def loaded_size(self):
        return self.size if self.model.llm is not None else 0

    def current_loaded_device(self):
        return self.load_device if self.model.llm is not None else self.offload_device

    def model_dtype(self):
        return None

    def model_patches_to(self, device):
        pass

    def model_patches_models(self):
        return []

    def is_dynamic(self):
        return False

    def is_clone(self, other):
        return other.model is self.model

    def lowvram_patch_counter(self):
        return 0

    def partially_load(self, device, extra_memory, force_patch_weights=False):
        if self.model.llm is not None:
            return 0
        chat_handler = get_chat_handler(self.model_file, self.mmproj_file) if self.mmproj_file else None
        self.model.llm = Llama(model_path=self.model_file, chat_handler=chat_handler, **self.llama_kwargs)
        return self.size

    def partially_unload(self, device, memory_to_free):
        return 0

    def detach(self, unpatch_all=True):
        if unpatch_all and self.model.llm is not None:
            self.model.llm.close()
            self.model.llm = None
        return self.model

class DavchaLLMLoader(io.ComfyNode):
    gguf_folder = os.path.join(folder_paths.models_dir, "LLM", "GGUF")
    
    # Cache pour les métadonnées (pour accélérer define_schema et F5)
    _metadata_cache = {}
    _path_map = {}
        
    @classmethod
    def define_schema(cls):
        path = os.path.join(cls.gguf_folder, "**", "*.gguf")
        files = list(filter(lambda x: "mmproj" not in x, glob(path, recursive=True)))
        
        # Nettoyage du cache si on a supprimé un fichier du dossier
        cls._metadata_cache = {k: v for k, v in cls._metadata_cache.items() if k in files}
        
        options = []
        cls._path_map.clear()
        
        for file in files:
            # On vérifie la date de modification du fichier
            mtime = os.path.getmtime(file)
            
            # On ne lit le GGUF que si c'est un nouveau fichier ou s'il a été modifié
            if file not in cls._metadata_cache or cls._metadata_cache[file]['mtime'] != mtime:
                try:
                    llm_meta = Llama(file, vocab_only=True, verbose=False)
                    arch = llm_meta.metadata.get("general.architecture", "llama")
                    max_n_ctx = int(llm_meta.metadata.get(f"{arch}.context_length", 8192))
                    max_n_gpu_layers = int(llm_meta.metadata.get(f"{arch}.block_count", 99))
                    
                    # Détection si c'est un modèle MoE (ex: DeepSeek, Mixtral, Qwen MoE, etc.)
                    expert_count = int(llm_meta.metadata.get(f"{arch}.expert_count", 0))
                    expert_used_count = int(llm_meta.metadata.get(f"{arch}.expert_used_count", 0))
                    is_moe = expert_count > 0 or expert_used_count > 0 or any(kw in arch.lower() for kw in ["moe", "mixtral", "dbrx"])
                    
                    cls._metadata_cache[file] = {
                        'mtime': mtime,
                        'max_n_ctx': max_n_ctx,
                        'max_n_gpu_layers': max_n_gpu_layers,
                        'is_moe': is_moe,
                        'expert_count': expert_count,
                        'expert_used_count': expert_used_count,
                    }
                except Exception as e:
                    print(f"[DavchaLLM] Erreur lecture métadonnées {file}: {e}")
                    cls._metadata_cache[file] = {'mtime': mtime, 'max_n_ctx': 8192, 'max_n_gpu_layers': 99, 'is_moe': False, 'expert_count': 0, 'expert_used_count': 0}
            
            cached_data = cls._metadata_cache[file]
            
            # Sub-inputs per selected model
            inputs = [
                io.Int.Input("n_ctx", min=512, max=cached_data['max_n_ctx'], default=4096),
                io.Int.Input("n_gpu_layers", min=-1, max=cached_data['max_n_gpu_layers'], default=-1),
            ]
            
            # Dynamically add n_cpu_moe ONLY for MoE models
            if cached_data['is_moe']:
                inputs.append(io.Int.Input("n_cpu_moe", min=0, max=cached_data['max_n_gpu_layers'], default=0))
                inputs.append(io.Int.Input("expert_used_count", min=1, max=cached_data['expert_count'], default=cached_data['expert_used_count']))

            name = Path(file).stem
            cls._path_map[name] = file
            option = io.DynamicCombo.Option(name, inputs)
            options.append(option)
        
        default_threads = max(1, multiprocessing.cpu_count() // 2)
        
        return io.Schema(
            node_id="DavchaLLMLoader",
            category="davcha/llm",
            inputs=[
                io.DynamicCombo.Input("model", options=options),
                io.Int.Input("n_batch", min=64, max=32768, default=2048),
                io.Int.Input("n_ubatch", min=64, max=4096, default=512),
                io.Int.Input("n_threads", min=1, max=128, default=default_threads),
                io.Combo.Input("kv_cache_type", options=["f16", "q8_0", "q4_0"], default="q8_0"),
                io.Boolean.Input("flash_attn", default=True),
            ],
            outputs=[
                io.Custom("DavchaLLMModel").Output(),
            ]
        )
    
    @classmethod
    def execute(cls, model, n_batch, n_ubatch, n_threads, kv_cache_type, flash_attn):
        n_ctx = model.get("n_ctx", 4096)
        n_gpu_layers = model.get("n_gpu_layers", -1)
        # Defaults to 0 safely for non-MoE models where "n_cpu_moe" is not present
        n_cpu_moe = model.get("n_cpu_moe", 0)
        expert_used_count = model.get("expert_used_count", None)
        
        model_name = model.get("model", None)
        if model_name is None:
            raise FileExistsError("model")
        
        model_file = cls._path_map[model_name]
        mmproj_files = glob(os.path.join(os.path.dirname(model_file), "*mmproj*.gguf"))
        mmproj = mmproj_files[0] if mmproj_files else None

        kv_type = KV_CACHE_TYPES.get(kv_cache_type, KV_CACHE_TYPES["q8_0"])
        
        try:
            meta_llm = Llama(model_path=model_file, vocab_only=True, verbose=False)
            arch = meta_llm.metadata.get("general.architecture", "llama")
        except:
            arch = "llama"
            
        kv_overrides = {}
        if expert_used_count is not None:
            override_key = f"{arch}.expert_used_count"
            kv_overrides[override_key] = int(expert_used_count)
            print(f"[DavchaLLM] Overriding MoE Top-K: {override_key} = {expert_used_count}")

        llama_kwargs = dict(
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            n_cpu_moe=n_cpu_moe,
            kv_overrides=kv_overrides,
            type_k=kv_type,
            type_v=kv_type,
            flash_attn=flash_attn,
            swa_full=True,
            n_batch=n_batch,
            n_ubatch=n_ubatch,
            n_threads=n_threads,
            n_threads_batch=n_threads,
            offload_kqv=True,
            logits_all=False,
            verbose=False
        )

        return io.NodeOutput(LlamaPatcher(model_file, mmproj, n_gpu_layers, llama_kwargs))

class DavchaLLM(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaLLM",
            category="davcha/llm",
            inputs=[
                io.Custom("DavchaLLMModel").Input("llm"),
                io.String.Input("system", multiline=True, dynamic_prompts=False),
                io.String.Input("prompt", multiline=True, dynamic_prompts=False),
                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                io.Int.Input("max_tokens", min=1, default=512),
                io.Float.Input("temperature", min=0.0, max=2.0, default=0.9),
                io.Float.Input("top_p", min=0.0, max=1.0, default=0.95),
                io.Int.Input("top_k", min=0, max=100, default=40),
                io.Float.Input("repeat_penalty", min=0.5, max=2.0, default=1.2),
                io.Boolean.Input("enable_thinking", default=True),
                io.String.Input("response_format", multiline=True, dynamic_prompts=False, default=""),
                io.Image.Input("images", optional=True),
            ],
            outputs=[
                io.String.Output()
            ]
        )
    
    @classmethod
    def execute(cls, llm, system, prompt, seed, max_tokens, temperature, top_p, top_k, repeat_penalty, enable_thinking, response_format, images=None):
        if images is not None:
            content = []
            if not isinstance(images, list):
                images = [images]
            for image in images:
                array = (image[0]*255).clamp(0, 255).byte().cpu().numpy()
                pil_img = Image.fromarray(array, mode="RGB")
                buf = BytesIO()
                pil_img.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                content.append({'type': 'image_url', 'image_url': {'url': f"data:image/png;base64,{b64}"}})
            
            content.append({'type': 'text', 'text': prompt})

            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": content}
            ]
        else:
            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt}
            ]

        response_format = json.loads(response_format) if response_format.strip() != "" else None

        llm = llm.load()

        with disable_thinking_context(llm, enable_thinking):
            result = llm.create_chat_completion(
                seed=seed,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repeat_penalty=repeat_penalty,
                response_format=response_format,
            )
        
        result = strip_via_fuzzy_tags(result['choices'][0]['message']['content'], llm)
        
        return io.NodeOutput(result)

class DavchaPromptEnricher(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaPromptEnricher",
            category="davcha/llm",
            inputs=[
                io.Custom("DavchaLLMModel").Input("llm"),
                io.String.Input("system", multiline=True, dynamic_prompts=False, default="You are a headless REST API server. You receive requests and output strictly valid JSON. You do not have a personality. You do not output markdown formatting. Output a single JSON object starting with {."),
                io.String.Input("user_instruction", multiline=True, dynamic_prompts=False, default="Here is a prompt with some variables:\n{prompt}\n\n---\n\nWrite detailed visual descriptions for the following variables: {keys}. Ensure the generated values can be replaced in the prompt, ensuring grammatical correctness and semantic coherence. Ensure no details from the variables are missing. Minimum length per key > 0. Result in the following JSON format: {json_format}"),
                io.String.Input("prompt", multiline=True, dynamic_prompts=False),
                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                io.Int.Input("max_tokens", min=1, default=512),
                io.Float.Input("temperature", min=0.0, max=2.0, default=0.9, step=0.01),
                io.Float.Input("top_p", min=0.0, max=1.0, default=0.95, step=0.01),
                io.Int.Input("top_k", min=0, max=100, default=40),
                io.Float.Input("repeat_penalty", min=0.5, max=2.0, default=1.2, step=0.01),
                io.Boolean.Input("enable_thinking", default=True),
                io.Boolean.Input("force_json_output", default=True),
                io.Image.Input("images", optional=True),
            ],
            outputs=[
                io.String.Output()
            ]
        )
    
    @classmethod
    def execute(cls, llm, system, user_instruction, prompt, seed, max_tokens, temperature, top_p, top_k, repeat_penalty, enable_thinking, force_json_output, images=None):
        if not re.search(r'\{([^}]+)\}', prompt):
            return io.NodeOutput(prompt)
        
        m = re.findall(r'\{([^}]+)\}', prompt)
        keys = {x: "" for x in m}
        json_format = {
            "type": "json_object",
            "schema": {
                "type": "object",
                "properties": {key:{"type":"string"} for key in m},
                "required": m
            }
        }

        p = user_instruction
        p = p.replace("{keys}", str(list(keys.keys())))
        p = p.replace("{json_format}", str(keys))
        p = p.replace("{prompt}", prompt)
        
        messages = [{"role": "system", "content": system or ""}] if system else []
        
        if images is not None:
            content = []
            if not isinstance(images, list):
                images = [images]
            for image in images:
                array = (image[0]*255).clamp(0, 255).byte().cpu().numpy()
                pil_img = Image.fromarray(array, mode="RGB")
                buf = BytesIO()
                pil_img.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                content.append({'type': 'image_url', 'image_url': {'url': f"data:image/png;base64,{b64}"}})
            
            content.append({'type': 'text', 'text': p})

            messages.append(
                {"role": "user", "content": content}
            )
        else:
            messages.append(
                {"role": "user", "content": p}
            )

        llm = llm.load()
        with disable_thinking_context(llm, enable_thinking):
            result = llm.create_chat_completion(
                seed=seed,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repeat_penalty=repeat_penalty,
                response_format=json_format if force_json_output else None,
            )
        groups = result['choices'][0]['message']['content']
        print(f'------------------\nLLM Response\n------------------\n\n{groups}\n\n------------------')
        groups = strip_via_fuzzy_tags(groups, llm)
        
        import json_repair
        p = prompt
        j = json_repair.loads(groups)
        d = {}
        if isinstance(j, list):
            for item in j:
                d.update(item)
        else:
            d = j
        for k, v in d.items():
            # find correct key
            best_match = process.extractOne(
                query=k,
                choices=list(keys.keys()),
                scorer=fuzz.partial_ratio
            )
            if best_match:
                matched_key, *_ = best_match
                p = p.replace(f"{{{matched_key}}}", "\n".join([str(x) for x in v]) if isinstance(v, list) else str(v))
            else:
                print(f"unmatched key: {k}")
        
        return io.NodeOutput(p)