import os
import re
import json
import tempfile
import zipfile
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

# ---- LLM: llama.cpp via llama_cpp_agent ----
from llama_cpp import Llama
from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
from llama_cpp_agent.providers import LlamaCppPythonProvider

# ----------------------
# Model configuration
# ----------------------
DEFAULT_REPO_ID = "tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF"
DEFAULT_FILENAME = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
MODELS_DIR = "models"
os.makedirs(MODELS_DIR, exist_ok=True)

def ensure_model(repo_id: str, filename: str) -> str:
    local_path = os.path.join(MODELS_DIR, filename)
    if not os.path.exists(local_path):
        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=MODELS_DIR)
    return local_path

# GPU context ensures model loads only when a GPU session is active
@spaces.GPU(duration=120)
def package_with_llm(ai_text, repo_id, filename, temperature, top_p, top_k, repeat_penalty, max_tokens, force_heuristic):
    model_path = ensure_model(repo_id, filename)
    llm = Llama(
        model_path=model_path,
        n_ctx=8192,
        n_gpu_layers=81,
        n_batch=1024,
        flash_attn=True,
    )
    provider = LlamaCppPythonProvider(llm)
    agent = LlamaCppAgent(
        provider,
        system_prompt=(
            "You are an expert code-packager and software project compiler. "
            "Given an AI-generated project description containing code blocks and hints "
            "about filenames or structure, extract each file with its most likely filename "
            "and exact code content. Return ONLY a strict JSON array named manifest, "
            "where each element is an object with keys 'filename' and 'content'. "
            "Do not add commentary outside JSON. "
            "Ensure filenames include directories if implied (e.g., 'src/main.py'). "
            "Preserve code exactly as provided inside code fences."
        ),
        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
        debug_output=False
    )

    prompt = f"""
Read the following AI project description and return ONLY JSON.
Output schema (strict):
[{{"filename": "server.js", "content": "// code..."}}]

AI project description:
{ai_text}
"""

    settings = provider.get_provider_default_settings()
    settings.temperature = float(temperature)
    settings.top_p = float(top_p)
    settings.top_k = int(top_k)
    settings.repeat_penalty = float(repeat_penalty)
    settings.max_tokens = int(max_tokens)
    settings.stream = False

    out = agent.get_chat_response(prompt, llm_sampling_settings=settings, print_output=False)

    # Parse JSON output robustly
    try:
        start, end = out.find('['), out.rfind(']')
        if start != -1 and end > start:
            manifest = json.loads(out[start:end+1])
        else:
            raise ValueError("No JSON found")
    except Exception:
        manifest = [{"filename": "project.txt", "content": out}]

    # Heuristic fallback if model fails
    if force_heuristic or (len(manifest) == 1 and manifest[0]["filename"] == "project.txt"):
        manifest = naive_regex_merge(ai_text)

    return create_zip_from_manifest(manifest)

def naive_regex_merge(text):
    blocks = []
    code_pattern = re.compile(r"```([a-zA-Z0-9]*)\n(.*?)```", re.DOTALL)
    lines = text.splitlines()
    candidates = []
    for line in lines:
        m = re.search(r"([A-Za-z0-9_\-./]+?\.[A-Za-z0-9]+)", line)
        if m:
            candidates.append(m.group(1))
    for idx, m in enumerate(code_pattern.finditer(text)):
        lang = m.group(1) or "txt"
        code = m.group(2)
        filename = candidates[idx] if idx < len(candidates) else f"file_{idx+1}.{lang}"
        blocks.append({"filename": filename, "content": code})
    return blocks

def create_zip_from_manifest(manifest):
    temp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(temp_dir, "project.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
        for item in manifest:
            fname = item.get("filename", "project.txt").lstrip("/")
            content = item.get("content", "")
            fpath = os.path.join(temp_dir, fname)
            os.makedirs(os.path.dirname(fpath), exist_ok=True)
            with open(fpath, "w", encoding="utf-8") as f:
                f.write(content)
            z.write(fpath, arcname=fname)
    return zip_path

# --------------------------
# Gradio Interface
# --------------------------
with gr.Blocks(title="AI Project Packager (GGUF, ZeroGPU)") as demo:
    gr.Markdown("# AI Project Packager (GGUF, ZeroGPU)")
    gr.Markdown("Uses ephemeral GPU power to infer filenames & package AI-generated code projects.")

    with gr.Row():
        ai_text = gr.Textbox(lines=24, label="Paste AI response here")

    with gr.Accordion("LLM Settings", open=False):
        repo_id = gr.Textbox(value=DEFAULT_REPO_ID, label="Model repo_id")
        filename = gr.Textbox(value=DEFAULT_FILENAME, label="Model filename (*.gguf)")
        with gr.Row():
            temperature = gr.Slider(0.0, 2.0, value=0.2, step=0.05, label="Temperature")
            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
            top_k = gr.Slider(0, 100, value=40, step=1, label="Top-k")
        with gr.Row():
            repeat_penalty = gr.Slider(0.8, 2.0, value=1.1, step=0.05, label="Repeat penalty")
            max_tokens = gr.Slider(256, 4096, value=2048, step=32, label="Max tokens")
        force_heuristic = gr.Checkbox(value=False, label="Force heuristic filename/code merge if JSON parse fails")

    out_zip = gr.File(label="Download packaged ZIP")

    run_btn = gr.Button("Package Project", variant="primary")
    run_btn.click(
        fn=package_with_llm,
        inputs=[ai_text, repo_id, filename, temperature, top_p, top_k, repeat_penalty, max_tokens, force_heuristic],
        outputs=[out_zip]
    )

if __name__ == "__main__":
    demo.launch()