import os
import sys
import subprocess
import site
import threading

# 1. BIND TO THE PERSISTENT COMPILATION REGISTRY
PERSISTENT_PACKAGES = "/data/compiled_cache"
os.environ["PYTHONUSERBASE"] = PERSISTENT_PACKAGES
os.environ["HF_HOME"] = "/data/.huggingface"
os.environ["XDG_CACHE_HOME"] = "/data/.cache"

TARGET_SITE_PATH = f"{PERSISTENT_PACKAGES}/lib/python3.12/site-packages"

if TARGET_SITE_PATH not in sys.path:
    sys.path.insert(0, TARGET_SITE_PATH)
    site.addsitedir(TARGET_SITE_PATH)

try:
    from llama_cpp import Llama
except ModuleNotFoundError:
    print("⏳ First boot detected. Compiling hardware engine into permanent storage /data...")
    os.makedirs(TARGET_SITE_PATH, exist_ok=True)
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", "--user",
        "llama-cpp-python", 
        "--extra-index-url", "https://github.io"
    ])
    site.addsitedir(TARGET_SITE_PATH)
    from llama_cpp import Llama

import gradio as gr
from huggingface_hub import hf_hub_download
import spaces

# 2. MODEL WEIGHT CONFIGURATIONS
print("Checking persistent storage for AI model weights...")
path_27b = hf_hub_download(repo_id="prism-ml/Bonsai-27B-gguf", filename="Bonsai-27B-Q1_0.gguf", local_dir="/data")
path_moe = hf_hub_download(repo_id="LiquidAI/LFM2-8B-A1B-GGUF", filename="LFM2-8B-A1B-Q4_K_M.gguf", local_dir="/data")

print("Initializing Liquid 8B MoE on active CPU thread...")
llm_cpu = Llama(model_path=path_moe, n_ctx=4096, n_gpu_layers=0, verbose=False)

# 3. NATIVE GRADIO WORKFLOWS (Ensures perfect ZeroGPU startup scans)
@spaces.GPU(duration=60)
def generate_27b(prompt):
    clean_prompt = str(prompt).strip()
    if not clean_prompt: return "Empty prompt."
    llm_gpu = Llama(model_path=path_27b, n_ctx=32768, n_gpu_layers=-1, n_batch=512, ctk_quant=2, ctv_quant=2, flash_attn=True, verbose=False)
    formatted = f"<|user|>\n{clean_prompt}<|endoftext|>\n<|assistant|>"
    response = llm_gpu(formatted, max_tokens=512, temperature=0.7)
    try: return response["choices"]["text"]
    except: return str(response)

def generate_moe_cpu(prompt):
    clean_prompt = str(prompt).strip()
    if not clean_prompt: return "Empty prompt."
    system_tool_prompt = "You are an advanced AI agent with Tool Calling capabilities."
    formatted = f"<|im_start|>system\n{system_tool_prompt}<|im_end|>\n<|im_start|>user\n{clean_prompt}<|im_end|>\n<|im_start|>assistant\n"
    response = llm_cpu(formatted, max_tokens=512, temperature=0.1)
    try: return response["choices"]["text"]
    except: return str(response)

# 4. START THE OPENAI COMPATIBILITY API AS A COMPANION PROCESS
def run_openai_api():
    try:
        import uvicorn
        # This triggers our proxy script securely in the background without locking ZeroGPU scans
        uvicorn.run("openai_api:api", host="0.0.0.0", port=8000, log_level="warning")
    except Exception as e:
        print(f"API Startup Error: {e}")

# Fire up the translator server on background thread port 8000
threading.Thread(target=run_openai_api, daemon=True).start()

# 5. GRADIO DASHBOARD
with gr.Blocks(title="Resilient AI Hub") as demo:
    gr.Markdown("# 🌳 Unstoppable Split-Brain AI Hub")
    with gr.Tab("🚀 Bonsai 27B (GPU Endpoint)"):
        input_27b = gr.Textbox(label="Enter prompt for 27B model (Uses Quota)", lines=6)
        output_27b = gr.Textbox(label="GPU Response Output")
        btn_27b = gr.Button("Submit to GPU")
        btn_27b.click(fn=generate_27b, inputs=input_27b, outputs=output_27b, api_name="chat")
    with gr.Tab("🪵 Liquid 8B MoE (CPU Endpoint / Native Tool Calling)"):
        input_moe = gr.Textbox(label="Enter prompt for MoE model (100% Free / Anti-Limit Backup)", lines=6)
        output_moe = gr.Textbox(label="MoE CPU Output")
        btn_moe = gr.Button("Submit to MoE Engine")
        btn_moe.click(fn=generate_moe_cpu, inputs=input_moe, outputs=output_moe, api_name="chat_backup")

if __name__ == "__main__":
    # Mount everything onto the default web configuration port
    demo.launch(server_name="0.0.0.0", server_port=7860)