"""
AIAGENT01 - Aapka Personal AI Agent
Built with smolagents + Gradio on Hugging Face Spaces
"""

import gradio as gr
from smolagents import CodeAgent, HfApiModel, tool, DuckDuckGoSearchTool
import subprocess, os, sys
from datetime import datetime

HF_TOKEN = os.environ.get("HF_TOKEN", None)

@tool
def execute_python(code: str) -> str:
    """Python code execute karo aur output dikhao.
    Args:
        code: Python code jo execute karni hai
    Returns:
        Output ya error
    """
    fname = f"/tmp/agent_code_{datetime.now().strftime('%Y%m%d_%H%M%S')}.py"
    with open(fname, "w") as f:
        f.write(code)
    try:
        r = subprocess.run(["python3", fname], capture_output=True, text=True, timeout=30)
        o = ""
        if r.stdout: o += f"OUTPUT:\n{r.stdout}\n"
        if r.stderr: o += f"ERRORS:\n{r.stderr}\n"
        if r.returncode != 0: o += f"Exit code: {r.returncode}"
        return o if o else "Done (no output)"
    except subprocess.TimeoutExpired:
        return "Timeout (30s)"
    except Exception as e:
        return f"Error: {str(e)}"
    finally:
        if os.path.exists(fname): os.remove(fname)

@tool
def install_package(pkg: str) -> str:
    """Install Python package via pip.
    Args:
        pkg: Package name
    Returns:
        Output
    """
    r = subprocess.run(["pip3", "install", pkg], capture_output=True, text=True, timeout=60)
    return (r.stdout + "\n" + r.stderr).strip() or "Installed"

@tool
def create_file(filename: str, content: str) -> str:
    """Create a new file with content.
    Args:
        filename: File path
        content: File content
    Returns:
        Success message
    """
    os.makedirs(os.path.dirname(os.path.abspath(filename)) or ".", exist_ok=True)
    with open(filename, "w") as f:
        f.write(content)
    return f"Created: {filename} ({len(content)} bytes)"

@tool
def read_file(filename: str) -> str:
    """Read file content.
    Args:
        filename: File path
    Returns:
        Content
    """
    if not os.path.exists(filename):
        return f"File not found: {filename}"
    with open(filename, "r") as f:
        return f.read()

@tool
def list_files(path: str = ".") -> str:
    """List files in directory.
    Args:
        path: Directory path
    Returns:
        File listing
    """
    try:
        files = os.listdir(path)
        if not files: return "(empty)"
        result = []
        for f in sorted(files):
            full = os.path.join(path, f)
            if os.path.isfile(full):
                result.append(f"  {f} ({os.path.getsize(full):,}b)")
            else:
                result.append(f"/ {f}/")
        return "\n".join(result)
    except Exception as e:
        return f"Error: {str(e)}"

@tool
def run_shell(command: str) -> str:
    """Run Linux shell command.
    Args:
        command: Command to run
    Returns:
        Output
    """
    try:
        r = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        o = ""
        if r.stdout: o += r.stdout
        if r.stderr: o += f"\n[stderr] {r.stderr}"
        return o if o else "Done"
    except subprocess.TimeoutExpired:
        return "Timeout (30s)"
    except Exception as e:
        return f"Error: {str(e)}"

# Agent setup
model = HfApiModel(model_id="Qwen/Qwen2.5-72B-Instruct", token=HF_TOKEN)

tools = [
    DuckDuckGoSearchTool(), execute_python, install_package,
    create_file, read_file, list_files, run_shell,
]

agent = CodeAgent(
    tools=tools, model=model, max_steps=20, verbose=True,
    additional_authorized_imports=[
        "requests", "json", "os", "subprocess", "datetime",
        "math", "random", "re", "collections", "itertools",
    ]
)

SYSTEM_PROMPT = "You are AIAGENT01. You can execute Python code, search web, create files, run commands. Always use tools, never just talk."

def chat_fn(message, history):
    if not message.strip():
        return "", history
    try:
        result = agent.run(f"{SYSTEM_PROMPT}
User: {message}")
        history.append((message, str(result)))
        return "", history
    except Exception as e:
        import traceback
        history.append((message, f"Error: {str(e)}"))
        return "", history

def clear_fn():
    return [], ""

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style="text-align:center;margin-bottom:15px">
        <h1 style="background:linear-gradient(135deg,#667eea,#764ba2);-webkit-background-clip:text;-webkit-text-fill-color:transparent;font-size:2.2em">
            AIAGENT01
        </h1>
        <p style="color:#666">Python code likho Web search karo Files banao Apps develop karo</p>
    </div>
    """)
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(height=450, bubble_full_width=False, avatar_images=(None, "🤖"))
            with gr.Row():
                msg = gr.Textbox(label="", placeholder="Yahan likho...", scale=5, container=False)
                btn = gr.Button("Send", variant="primary", scale=1)
            with gr.Row():
                clr = gr.Button("Clear", variant="secondary", size="sm")
        with gr.Column(scale=1):
            gr.Markdown("""
            **Example Commands:**
            1. "Python mein calculator banao aur run karo"
            2. "Aaj ki news search karo"
            3. "HTML website banao"
            4. "System info dikhao"
            5. "1-50 Fibonacci print karo"
            """)
    msg.submit(chat_fn, [msg, chatbot], [msg, chatbot])
    btn.click(chat_fn, [msg, chatbot], [msg, chatbot])
    clr.click(clear_fn, None, [chatbot, msg])

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