import os
import subprocess
import gradio as ui
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Cache directory explicitly set kar rahe hain taaki permission error na aaye
os.environ["HF_HOME"] = "/tmp/hf_cache"

# Ultra-Lightweight Uncensored Base Model (Free CPU ke liye ekdum perfect aur fast)
MODEL_REPO = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
MODEL_FILE = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"

print(f"Downloading model: {MODEL_FILE}...")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, cache_dir="/tmp/hf_cache")
print("Model downloaded successfully!")

# Initialize Llama model (Optimized context and threads for stable CPU deployment)
llm = Llama(model_path=model_path, n_ctx=1024, n_threads=4)

# Function to execute Terminal Commands securely
def run_terminal_command(command):
    try:
        result = subprocess.run(
            command, 
            shell=True, 
            text=True, 
            capture_output=True, 
            timeout=30
        )
        
        output = ""
        if result.stdout:
            output += f"--- Terminal Output ---\n{result.stdout}\n"
        if result.stderr:
            output += f"--- Terminal Error ---\n{result.stderr}\n"
        if not output:
            output = "Command executed successfully (No output)."
            
        return output
    except Exception as e:
        return f"Error executing command: {str(e)}"

# Core API and Chat function
def api_predict(message):
    system_prompt = (
        "You are an AI assistant with terminal access. "
        "If the user wants to run a shell/terminal command, reply ONLY with the exact command prefixed by 'RUN_CMD: '. "
        "For example, if asked to list files, reply with 'RUN_CMD: ls'. "
        "If the user is just chatting, reply with a normal conversational response."
    )
    
    prompt = f"System: {system_prompt}\nUser: {message}\nAssistant:"
    
    response = llm(
        prompt,
        max_tokens=256,
        stop=["User:", "\n"],
        echo=False
    )
    
    ai_output = response['choices'][0]['text'].strip()
    
    # Check if AI triggered a terminal command
    if ai_output.startswith("RUN_CMD:"):
        cmd_to_run = ai_output.replace("RUN_CMD:", "").strip()
        terminal_result = run_terminal_command(cmd_to_run)
        return f"⚠️ [COMMAND EXECUTION]\nCommand: {cmd_to_run}\n\n{terminal_result}"
    
    return ai_output

# Gradio Interface Setup
demo = ui.Interface(
    fn=api_predict,
    inputs=ui.Textbox(label="Input Message / Command"),
    outputs=ui.Textbox(label="Response"),
    title="Lightweight Uncensored API Server"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)