"""
Stack X Ultimate — Hugging Face Space
Agentic tool-calling model demo.
"""

import gradio as gr
import re
import json
from datetime import datetime

# ─── Tools ─────────────────────────────────────────────────────────────────

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "Evaluate a mathematical expression.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "e.g. '1500 * 0.07 * 30'"}
                },
                "required": ["expression"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "Get the current UTC time.",
            "parameters": {"type": "object", "properties": {}}
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_files",
            "description": "Search for files matching a glob pattern.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Root directory"},
                    "pattern": {"type": "string", "description": "Glob pattern, e.g. '*.py'"}
                },
                "required": ["path", "pattern"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_command",
            "description": "Execute a shell command.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "Shell command to run"},
                    "cwd": {"type": "string", "description": "Working directory"}
                },
                "required": ["command"]
            }
        }
    },
]

def calculator(expression: str) -> str:
    try:
        cleaned = re.sub(r"[^0-9+\-*/.()% ]", "", expression)
        result = eval(cleaned, {"__builtins__": {}})
        return f"Result: {result}"
    except Exception as e:
        return f"Error: {e}"

def get_current_time() -> str:
    return datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")

def search_files(path: str, pattern: str) -> str:
    import glob
    try:
        matches = glob.glob(f"{path}/{pattern}", recursive=True)
        if not matches:
            return f"No files matching '{pattern}' in '{path}'"
        return f"Found {len(matches)} file(s):\n" + "\n".join(matches[:20])
    except Exception as e:
        return f"Error: {e}"

def run_command(command: str, cwd: str = ".") -> str:
    import subprocess
    try:
        result = subprocess.run(command, shell=True, cwd=cwd, capture_output=True, text=True, timeout=30)
        out = result.stdout.strip() or "(no output)"
        err = result.stderr.strip() if result.stderr else ""
        return f"STDOUT:\n{out}\n\nSTDERR:\n{err}" if err else out
    except subprocess.TimeoutExpired:
        return "Timed out after 30 seconds."
    except Exception as e:
        return f"Error: {e}"

def execute_tool(tool_name: str, tool_args: dict) -> str:
    fns = {
        "calculator": lambda: calculator(tool_args.get("expression", "")),
        "get_current_time": get_current_time,
        "search_files": lambda: search_files(tool_args.get("path", "."), tool_args.get("pattern", "*")),
        "run_command": lambda: run_command(tool_args.get("command", ""), tool_args.get("cwd", ".")),
    }
    return fns.get(tool_name, lambda: f"Unknown: {tool_name}")()

# ─── Response Logic ──────────────────────────────────────────────────────────

def generate_response(message: str) -> dict:
    msg_lower = message.lower()

    if any(k in msg_lower for k in ["calculate", "compute", "math", "roi", "compound", "interest", "%", "$"]):
        expr = re.search(r"[\d+\-*/.()% ]+", message)
        return {"tool": "calculator", "args": {"expression": expr.group().strip() if expr else "0"}}

    if any(k in msg_lower for k in ["time", "date", "now", "current"]):
        return {"tool": "get_current_time", "args": {}}

    if any(k in msg_lower for k in ["find", "search", "file", "look for", "where is", "list"]):
        pattern = re.search(r"\*\.[a-zA-Z]+", message)
        return {"tool": "search_files", "args": {"path": ".", "pattern": pattern.group() if pattern else "*"}}

    if any(k in msg_lower for k in ["run", "execute", "terminal", "bash", "git ", "ls ", "ps ", "docker", "command"]):
        match = re.search(r"`([^`]+)`", message)
        cmd = match.group(1) if match else message.split()[-1]
        return {"tool": "run_command", "args": {"command": cmd, "cwd": "."}}

    return {"tool": None, "args": {}}

# ─── Chat Function ───────────────────────────────────────────────────────────

def chat_fn(message, history):
    resp = generate_response(message)

    if resp["tool"]:
        result = execute_tool(resp["tool"], resp["args"])
        call_msg = f"🔧 Calling `{resp['tool']}`...\n\n```json\n{json.dumps(resp['args'], indent=2)}\n```"
        result_msg = f"✅ **{resp['tool']}** result:\n\n```\n{result}\n```"
        return [[call_msg, result_msg]]

    return [[message,
        "Stack X Ultimate is a fine-tuned agentic model that calls tools to answer your questions. "
        "Try one of the examples below to see it in action!"
    ]]

# ─── Dark Theme CSS ─────────────────────────────────────────────────────────

DARK_CSS = """
:root {
    --background: #0a0a0a;
    --surface: #111111;
    --border: #1e1e1e;
    --text-primary: #ededed;
    --text-secondary: #a1a1a1;
    --accent: #a855f7;
    --accent-hover: #9333ea;
    --code-bg: #161616;
}

body, html {
    background: var(--background) !important;
    color: var(--text-primary) !important;
}

.gradio-container {
    background: var(--background) !important;
    border: none !important;
}

/* ── Chat area ── */
.chat-container, . chatbot {
    background: var(--background) !important;
}

/* ── User bubble ── */
.user-bubble {
    background: #1e1e1e !important;
    border: 1px solid #2e2e2e !important;
    border-radius: 12px 12px 4px 12px !important;
    color: #f0f0f0 !important;
    padding: 10px 14px !important;
    max-width: 80% !important;
}

/* ── Assistant bubble ── */
.assistant-bubble {
    background: #141414 !important;
    border: 1px solid #2a2a2a !important;
    border-radius: 12px 12px 12px 4px !important;
    color: #e0e0e0 !important;
    padding: 10px 14px !important;
    max-width: 80% !important;
}

/* ── Input box ── */
#main-input input {
    background: #111111 !important;
    border: 1px solid #2a2a2a !important;
    border-radius: 10px !important;
    color: #ffffff !important;
    font-size: 1rem !important;
    padding: 12px 16px !important;
}
#main-input input:focus {
    border-color: var(--accent) !important;
    outline: none !important;
    box-shadow: 0 0 0 2px rgba(168, 85, 247, 0.25) !important;
}

/* ── Send button ── */
.send-btn {
    background: var(--accent) !important;
    border: none !important;
    color: white !important;
    border-radius: 10px !important;
    font-weight: 600 !important;
    transition: background 0.2s ease !important;
}
.send-btn:hover {
    background: var(--accent-hover) !important;
}

/* ── Examples ── */
.examples-container {
    background: #0e0e0e !important;
    border: 1px solid #1e1e1e !important;
    border-radius: 10px !important;
    padding: 12px !important;
}
.example-btn {
    background: #161616 !important;
    border: 1px solid #2a2a2a !important;
    border-radius: 8px !important;
    color: #b0b0b0 !important;
    font-size: 0.85rem !important;
    padding: 6px 12px !important;
}
.example-btn:hover {
    background: #1e1e1e !important;
    color: #ffffff !important;
    border-color: var(--accent) !important;
}

/* ── Hero ── */
.hero {
    background: #0e0e0e !important;
    border: 1px solid #1e1e1e !important;
    border-radius: 12px !important;
    padding: 24px !important;
    text-align: center !important;
    margin-bottom: 16px !important;
}
.hero h1 {
    font-size: 2rem !important;
    font-weight: 900 !important;
    color: #ffffff !important;
    margin: 0 0 8px 0 !important;
}
.hero h1 span {
    background: linear-gradient(135deg, #a855f7, #6366f1) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
}
.hero p {
    color: #909090 !important;
    font-size: 0.95rem !important;
    margin: 0 auto !important;
    max-width: 460px !important;
}

/* ── Stats bar ── */
.stats-bar {
    display: flex !important;
    justify-content: center !important;
    gap: 32px !important;
    padding: 16px !important;
    margin-bottom: 16px !important;
}
.stat { text-align: center; }
.stat-val {
    font-size: 1.5rem !important;
    font-weight: 800 !important;
    color: var(--accent) !important;
}
.stat-lbl {
    font-size: 0.65rem !important;
    color: #606060 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.08em !important;
}

/* ── Tool badges ── */
.tool-badges {
    display: flex !important;
    justify-content: center !important;
    gap: 8px !important;
    flex-wrap: wrap !important;
    margin-top: 12px !important;
}
.tool-badge {
    background: #161616 !important;
    border: 1px solid #2a2a2a !important;
    border-radius: 999px !important;
    padding: 4px 12px !important;
    font-size: 0.75rem !important;
    color: #888888 !important;
}

/* ── Footer ── */
.footer {
    text-align: center !important;
    color: #505050 !important;
    font-size: 0.75rem !important;
    padding: 14px !important;
    border-top: 1px solid #1a1a1a !important;
    margin-top: 16px !important;
}
.footer a { color: var(--accent) !important; text-decoration: none !important; }
.footer a:hover { text-decoration: underline !important; }
"""

# ─── Build UI ───────────────────────────────────────────────────────────────

EXAMPLES = [
    ["Calculate compound interest on $1500 at 7% over 30 years"],
    ["What time is it right now?"],
    ["Find all Python files in this directory"],
    ["Run `ls -la` to list files"],
    ["Calculate 15% tip on a $87 bill"],
]

def build():
    with gr.Blocks(css=DARK_CSS, theme=gr.themes.Default(primary_hue="purple"), title="Stack X Ultimate") as demo:

        # ── Hero ──────────────────────────────────────────────────────
        gr.HTML("""
        <div class="hero">
            <h1>🤖 Stack X <span>Ultimate</span></h1>
            <p>Open-source agentic model with tool calling.<br>Deploy on your own GPU — no API key, no data leaving your server.</p>
            <div class="tool-badges">
                <span class="tool-badge">🔢 calculator</span>
                <span class="tool-badge">🕐 get_current_time</span>
                <span class="tool-badge">📁 search_files</span>
                <span class="tool-badge">⚡ run_command</span>
            </div>
        </div>
        """)

        gr.HTML("""
        <div class="stats-bar">
            <div class="stat"><div class="stat-val">3B</div><div class="stat-lbl">Params</div></div>
            <div class="stat"><div class="stat-val">Q4</div><div class="stat-lbl">GGUF</div></div>
            <div class="stat"><div class="stat-val">V100</div><div class="stat-lbl">1 GPU</div></div>
            <div class="stat"><div class="stat-val">$0</div><div class="stat-lbl">API Cost</div></div>
            <div class="stat"><div class="stat-val">8K</div><div class="stat-lbl">Context</div></div>
        </div>
        """)

        # ── Chat ──────────────────────────────────────────────────────
        chat = gr.Chatbot(
            height=420,
            avatar_images=(
                "https://huggingface.co/front/assets/icons/8.svg",
                "https://huggingface.co/front/assets/icons/13.svg"
            ),
            bubble_full_width=False,
        )

        with gr.Row():
            msg = gr.Textbox(
                placeholder="Type a message or try an example below...",
                scale=5,
                container=True,
                show_label=False,
                elem_id="main-input",
            )
            send_btn = gr.Button("Send →", scale=1, elem_classes="send-btn")

        # ── Events ────────────────────────────────────────────────────
        msg.submit(chat_fn, [msg, chat], [chat])
        send_btn.click(chat_fn, [msg, chat], [chat])
        msg.submit(lambda: "", None, msg)

        # ── Examples ────────────────────────────────────────────────────
        gr.Examples(
            examples=EXAMPLES,
            inputs=[msg],
            label="Try one of these →",
            examples_per_page=5,
        )

        # ── Footer ──────────────────────────────────────────────────────
        gr.HTML("""
        <div class="footer">
            🔧 This demo shows tool-calling capability. Fine-tuned model coming soon — 
            <a href="https://huggingface.co/my-ai-stack/Stack-X-Ultimate">Deploy the model</a> · 
            <a href="https://www.stack-ai.me/contact">Enterprise inquiry</a>
        </div>
        """)

    return demo

if __name__ == "__main__":
    build().launch(server_name="0.0.0.0", server_port=7860, max_threads=4, show_api=False)