import gradio as gr
import subprocess
import os
import json
import requests

# ============================================
# TOOL DEFINITIONS
# ============================================

def tool_bash(command: str) -> str:
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return result.stdout[:3000] if result.stdout else result.stderr[:3000] or "Done"
    except Exception as e:
        return f"Error: {str(e)}"

def tool_write_file(path: str, content: str) -> str:
    try:
        with open(path, 'w') as f:
            f.write(content)
        return f"Written to {path}"
    except Exception as e:
        return f"Error: {str(e)}"

def tool_read_file(path: str) -> str:
    try:
        with open(path, 'r') as f:
            return f.read()[:5000]
    except Exception as e:
        return f"Error: {str(e)}"

def tool_list_files(directory: str = ".") -> str:
    try:
        return "\n".join(os.listdir(directory)[:50])
    except Exception as e:
        return f"Error: {str(e)}"

TOOLS = {
    "bash": tool_bash,
    "write_file": tool_write_file,
    "read_file": tool_read_file,
    "list_files": tool_list_files,
}

# ============================================
# AI LOGIC
# ============================================

from huggingface_hub import InferenceClient
import os

# Get token from environment (set in Space secrets)
HF_TOKEN = os.environ.get("HF_TOKEN")

# Try multiple models
client = InferenceClient(token=HF_TOKEN)

def get_response(messages):
    """Get response from AI model"""
    try:
        # Use free serverless model
        prompt = ""
        for m in messages:
            if m["role"] == "system":
                prompt += f"System: {m['content']}\n"
            elif m["role"] == "user":
                prompt += f"User: {m['content']}\n"
            elif m["role"] == "assistant":
                prompt += f"Assistant: {m['content']}\n"
        prompt += "Assistant:"
        
        # Use free inference endpoint
        import requests
        
        API_URL = "https://api-inference.huggingface.co/models/microsoft/DialoGPT-medium"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
        
        response = requests.post(
            API_URL,
            headers=headers,
            json={"inputs": prompt[-1000:], "parameters": {"max_new_tokens": 256}}
        )
        
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and len(result) > 0:
                return result[0].get("generated_text", "No response").split("Assistant:")[-1].strip()
            return str(result)
        else:
            return f"API Error: {response.status_code} - {response.text[:200]}"
            
    except Exception as e:
        return f"Error: {str(e)}"

SYSTEM = """You are AI7, a helpful AI assistant.

Tools available (use JSON format to call):
- bash: {"tool": "bash", "command": "ls"}
- write_file: {"tool": "write_file", "path": "file.txt", "content": "text"}
- read_file: {"tool": "read_file", "path": "file.txt"}
- list_files: {"tool": "list_files", "dir": "."}

When you need to use a tool, respond with ONLY the JSON, nothing else.
Otherwise, respond normally to help the user."""

def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM}]
    
    for user, bot in history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})
    
    messages.append({"role": "user", "content": message})
    
    reply = get_response(messages)
        
    # Check for tool call
    try:
        tool_data = json.loads(reply.strip())
        if "tool" in tool_data:
            tool_name = tool_data["tool"]
            if tool_name == "bash":
                result = tool_bash(tool_data.get("command", ""))
            elif tool_name == "write_file":
                result = tool_write_file(tool_data.get("path", ""), tool_data.get("content", ""))
            elif tool_name == "read_file":
                result = tool_read_file(tool_data.get("path", ""))
            elif tool_name == "list_files":
                result = tool_list_files(tool_data.get("dir", "."))
            else:
                result = "Unknown tool"
            
            # Get final response with tool result
            messages.append({"role": "assistant", "content": reply})
            messages.append({"role": "user", "content": f"Tool result:\n{result}"})
            reply = get_response(messages)
    except:
        pass
    
    return reply

# ============================================
# INTERFACE
# ============================================

with gr.Blocks(title="AI7") as demo:
    gr.Markdown("# AI7 - AI Agent\n\nAsk me anything. I can run commands and manage files.")
    
    with gr.Row():
        inp = gr.Textbox(label="Message", scale=4)
        btn = gr.Button("Send", variant="primary")
    
    out = gr.Textbox(label="Response", lines=10)
    
    btn.click(chat, [inp, gr.State([])], out)
    inp.submit(chat, [inp, gr.State([])], out)

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