import os
import subprocess
import shutil
from huggingface_hub import hf_hub_download
from typing import Dict, Any

class EndpointHandler:
    def __init__(self, path: str):
        # 1. Download GGUF model dynamically
        print("Downloading DiffusionGemma Q4_K_M GGUF model...")
        self.model_path = hf_hub_download(
            repo_id="unsloth/diffusiongemma-26B-A4B-it-GGUF",
            filename="diffusiongemma-26B-A4B-it-Q4_K_M.gguf"
        )
        print(f"Model downloaded to: {self.model_path}")

        # 2. Clone and build llama.cpp with PR 24423 (diffusiongemma support)
        print("Cloning llama.cpp...")
        if os.path.exists("/tmp/llama.cpp"):
            shutil.rmtree("/tmp/llama.cpp")
            
        subprocess.run("git clone https://github.com/ggml-org/llama.cpp.git /tmp/llama.cpp", shell=True, check=True)
        
        print("Checking out PR 24423 (diffusiongemma branch)...")
        subprocess.run("cd /tmp/llama.cpp && git fetch origin pull/24423/head:diffusiongemma && git checkout diffusiongemma", shell=True, check=True)
        
        # Check if GPU exists
        has_cuda = os.path.exists("/usr/local/cuda") or subprocess.run("nvidia-smi", shell=True, capture_output=True).returncode == 0
        
        print(f"Building llama-diffusion-cli (CUDA={has_cuda})...")
        if has_cuda:
            build_cmd = "cd /tmp/llama.cpp && cmake -B build -DGGML_CUDA=ON && cmake --build build --target llama-diffusion-cli -j 4"
        else:
            build_cmd = "cd /tmp/llama.cpp && cmake -B build && cmake --build build --target llama-diffusion-cli -j 4"
            
        subprocess.run(build_cmd, shell=True, check=True)
        self.binary_path = "/tmp/llama.cpp/build/bin/llama-diffusion-cli"
        print("Successfully compiled llama-diffusion-cli!")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.get("inputs", "")
        if not inputs:
            return {"error": "Missing 'inputs' field in payload"}
            
        # Check if GPU is present to set ngl
        has_cuda = os.path.exists("/usr/local/cuda") or os.path.exists("/dev/nvidia0")
        ngl = 99 if has_cuda else 0
        
        # Execute llama-diffusion-cli
        cmd = [self.binary_path, "-m", self.model_path, "-p", inputs, "-ngl", str(ngl), "-n", "300"]
        print(f"Executing: {' '.join(cmd)}")
        
        try:
            res = subprocess.run(cmd, capture_output=True, text=True, timeout=120, check=True)
            return {"generated_text": res.stdout}
        except subprocess.CalledProcessError as e:
            return {"error": f"Execution failed: {e.stderr}\nOutput: {e.stdout}"}
        except Exception as e:
            return {"error": str(e)}