import os
import zipfile
import shutil
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Fast download enabled
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
CACHE_DIR = os.path.join(os.getcwd(), "model_cache")
SANDBOX_DIR = os.path.join(os.getcwd(), "sandbox")

# Ensure sandbox directory exists
if not os.path.exists(SANDBOX_DIR):
    os.makedirs(SANDBOX_DIR)

llm = None

def load_model():
    global llm
    print("--- 📥 Downloading 26B Model ---")
    model_path = hf_hub_download(
        repo_id="BugTraceAI/BugTraceAI-Apex-G4-26B-Q4",
        filename="BugTraceAI-Apex-G4-26B-Q4.gguf",
        cache_dir=CACHE_DIR
    )
    llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2, use_mmap=True)
    return "Model Ready!"

def handle_zip(file_obj):
    if file_obj is None:
        return "No file uploaded."
    
    # Clear previous sandbox data
    if os.path.exists(SANDBOX_DIR):
        shutil.rmtree(SANDBOX_DIR)
    os.makedirs(SANDBOX_DIR)

    try:
        if file_obj.name.endswith(".zip"):
            with zipfile.ZipFile(file_obj.name, 'r') as zip_ref:
                zip_ref.extractall(SANDBOX_DIR)
            
            # List files for the AI context
            files = []
            for root, dirs, filenames in os.walk(SANDBOX_DIR):
                for f in filenames:
                    files.append(os.path.relpath(os.path.join(root, f), SANDBOX_DIR))
            
            file_list_str = "\n".join(files[:20]) # Limit to first 20 files to save tokens
            return f"ZIP Extracted. Files found:\n{file_list_str}"
        else:
            return "Please upload a .zip file."
    except Exception as e:
        return f"Error: {e}"

def chat_func(message, history, zip_file):
    if llm is None:
        return "Model is still loading... wait ~5 mins."
    
    # Logic to tell the AI what is in the sandbox
    context = ""
    if zip_file:
        files_in_sandbox = [f for f in os.listdir(SANDBOX_DIR)]
        context = f"The user has uploaded a ZIP. Files available in /sandbox: {files_in_sandbox}\n"
    
    prompt = f"{context}User: {message}\nAssistant:"
    output = llm(prompt, max_tokens=512, stop=["User:"])
    return output["choices"][0]["text"]

with gr.Blocks() as demo:
    gr.Markdown("# 📦 ZIP-Enabled AI Sandbox")
    
    with gr.Row():
        zip_input = gr.File(label="Upload ZIP Project", file_types=[".zip"])
        status_output = gr.Textbox(label="Extraction Status")
    
    zip_input.change(handle_zip, inputs=zip_input, outputs=status_output)
    
    gr.ChatInterface(fn=chat_func, additional_inputs=[zip_input])

if __name__ == "__main__":
    import threading
    threading.Thread(target=load_model).start()
    demo.launch(server_name="0.0.0.0", server_port=7860)