import gradio as gr
import json
import re
import math
import os
import time
from dataclasses import dataclass, field
from typing import List, Dict, Optional
from enum import Enum

# ============================================================
# Monkey-patch Gradio 5.9.0 get_api_info crash
# Fixes: TypeError: argument of type 'bool' is not iterable
# in gradio_client/utils.py:887 (if "const" in schema:)
# ============================================================
import gradio_client.utils as gradio_utils
_orig_get_type = gradio_utils.get_type
def _patched_get_type(schema):
    if isinstance(schema, bool):
        return "boolean"
    return _orig_get_type(schema)
gradio_utils.get_type = _patched_get_type

# ============================================================
# Tool Definitions
# ============================================================

def safe_calc(expression: str) -> str:
    """Safe calculator using restricted eval."""
    import ast, operator as op
    ops = {
        ast.Add: op.add, ast.Sub: op.sub, ast.Mult: op.mul,
        ast.Div: op.truediv, ast.Pow: op.pow, ast.Mod: op.mod,
        ast.FloorDiv: op.floordiv, ast.USub: op.neg, ast.UAdd: op.pos,
    }
    funcs = {
        "abs": abs, "round": round, "int": int, "float": float,
        "min": min, "max": max, "sum": sum, "len": len,
        "sqrt": math.sqrt, "log": math.log, "sin": math.sin, "cos": math.cos,
        "pi": lambda: math.pi, "e": lambda: math.e,
    }
    consts = {"pi": math.pi, "e": math.e, "tau": math.tau}
    blocked = ["__", "import", "exec", "eval", "open", "os.", "subprocess", "sys."]
    
    for b in blocked:
        if b in expression.lower():
            return f"❌ Blocked pattern: {b}"
    
    def eval_node(node):
        if isinstance(node, ast.Expression):
            return eval_node(node.body)
        elif isinstance(node, ast.Constant):
            return node.value
        elif isinstance(node, ast.Num):
            return node.n
        elif isinstance(node, ast.Name):
            if node.id in consts:
                return consts[node.id]
            raise NameError(f"Unknown: {node.id}")
        elif isinstance(node, ast.UnaryOp):
            return ops[type(node.op)](eval_node(node.operand))
        elif isinstance(node, ast.BinOp):
            return ops[type(node.op)](eval_node(node.left), eval_node(node.right))
        elif isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name) and node.func.id in funcs:
                args = [eval_node(a) for a in node.args]
                return funcs[node.func.id](*args) if callable(funcs[node.func.id]) else funcs[node.func.id]()
            raise NameError(f"Unknown function")
        raise ValueError(f"Unsupported: {type(node).__name__}")
    
    try:
        tree = ast.parse(expression.strip(), mode='eval')
        result = eval_node(tree.body)
        if isinstance(result, float):
            if result == int(result) and abs(result) < 1e15:
                return str(int(result))
            return f"{result:.4f}".rstrip("0").rstrip(".")
        return str(result)
    except Exception as e:
        return f"❌ Error: {e}"


MOCK_KB = {
    "python": "Python is a high-level, general-purpose programming language created by Guido van Rossum in 1991.",
    "langgraph": "LangGraph is a library for building stateful, multi-actor applications with LLMs. It extends LangChain by adding graph-based orchestration of agents.",
    "langchain": "LangChain is a framework for developing applications powered by language models. It provides tools, chains, and agents.",
    "qwen": "Qwen2.5 is Alibaba Cloud's LLM series. Qwen2.5-1.5B has 1.5B parameters and supports 32K context.",
    "gradio": "Gradio is a Python library for building ML web demos. It provides UI components for models.",
    "huggingface": "Hugging Face provides the Transformers library, model hub, and Spaces for ML demos.",
    "machine learning": "Machine learning enables systems to learn patterns from data without being explicitly programmed.",
    "transformer": "A deep learning architecture using self-attention, introduced in 'Attention Is All You Need' (2017).",
    "kaggle": "Kaggle is a data science community platform owned by Google, offering competitions, datasets, and notebooks.",
    "agent": "An AI agent perceives its environment, makes decisions, and takes actions to achieve goals. Tool-calling agents use external tools.",
}

def web_search(query: str) -> str:
    """Mock web search with knowledge base."""
    q = query.lower().strip()
    results = []
    for kw, info in MOCK_KB.items():
        if kw in q:
            results.append(f"📖 **{kw.title()}**: {info}")
    if results:
        return "\n\n".join(results[:3])
    return f"📭 No results found for '{query}'. Try rephrasing."


def read_file(path: str) -> str:
    """Read a text file safely."""
    if ".." in path:
        return "❌ Directory traversal blocked."
    try:
        if not os.path.exists(path):
            return f"❌ File not found: {path}"
        if os.path.isdir(path):
            items = os.listdir(path)
            return f"📁 Directory: {path}\n" + "\n".join(f"  {'📄' if os.path.isfile(os.path.join(path,f)) else '📁'} {f}" for f in items[:20])
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read(2000)
        return f"📄 {path}\n\n{content}"
    except Exception as e:
        return f"❌ Error: {e}"


TOOLS = {
    "calculator": {"fn": safe_calc, "desc": "Evaluate math expressions (e.g., 2 + 2, sqrt(144), pi * 5^2)"},
    "web_search": {"fn": web_search, "desc": "Search the knowledge base for information"},
    "file_reader": {"fn": read_file, "desc": "Read a text file from the filesystem"},
}

# ============================================================
# Agent Engine
# ============================================================

class StepType(Enum):
    THOUGHT = "thought"
    TOOL_CALL = "tool_call" 
    TOOL_RESULT = "tool_result"
    FINAL = "final"

@dataclass
class AgentStep:
    type: StepType
    content: str
    tool_name: Optional[str] = None
    tool_input: Optional[str] = None
    tool_output: Optional[str] = None
    duration_ms: float = 0.0

def detect_intent(query: str) -> dict:
    """Detect what the user wants and route to appropriate tool."""
    q = query.lower().strip()
    
    # Calculator patterns
    calc_patterns = [
        r"(?:calculate|compute|what\s+is|evaluate|solve|how\s+much\s+is)\s+(.+)",
        r"(.+)\s*[+\-*/^%].+",  # contains math operators
    ]
    for pat in calc_patterns:
        m = re.search(pat, q)
        if m:
            expr = m.group(1) if m.lastindex else q
            # Clean up the expression
            expr = re.sub(r"^(?:calculate|compute|what\s+is|evaluate|solve|how\s+much\s+is)\s+", "", expr, flags=re.IGNORECASE).strip()
            if any(op in expr for op in ["+", "-", "*", "/", "^", "%", "sqrt", "log", "sin", "cos", "abs", "round", "pi", "e"]):
                return {"tool": "calculator", "input": expr}
    
    # File reader patterns
    if re.search(r"(?:read|open|show|list|cat|view|display|contents of)\s+(?:file\s+)?(.+)", q):
        m = re.search(r"(?:read|open|show|list|cat|view|display|contents of)\s+(?:file\s+)?(.+)", q)
        path = m.group(1).strip().strip('"\'')
        return {"tool": "file_reader", "input": path}
    
    if q.startswith("read ") or q.startswith("open ") or q.startswith("list "):
        parts = q.split(" ", 1)
        if len(parts) > 1:
            return {"tool": "file_reader", "input": parts[1].strip()}
    
    # Web search patterns (everything else with a question)
    if any(w in q for w in ["what", "who", "when", "where", "why", "how", "tell me", "explain", "about", "define"]):
        return {"tool": "web_search", "input": q}
    
    # Default: check for math operators
    if re.search(r"[\d\s]*[+\-*/^][\d\s]*", q):
        return {"tool": "calculator", "input": q}
    
    # Greetings - no tool needed
    greetings = ["hi", "hello", "hey", "greetings", "good morning", "good afternoon", "good evening"]
    if any(g in q for g in greetings):
        return {"tool": None, "input": q}
    
    # Fallback to web search
    return {"tool": "web_search", "input": q}


def run_agent(query: str) -> List[AgentStep]:
    """Run the agent pipeline and return all steps."""
    steps = []
    t_start = time.time()
    
    # Step 1: Thought
    thought_start = time.time()
    intent = detect_intent(query)
    thought_duration = (time.time() - thought_start) * 1000
    
    if intent["tool"] is None:
        # Direct response (no tool needed)
        steps.append(AgentStep(
            type=StepType.THOUGHT,
            content=f"The user said: '{query}'. This appears to be a greeting or simple statement — no tool needed.",
            duration_ms=thought_duration,
        ))
        steps.append(AgentStep(
            type=StepType.FINAL,
            content=f"Hello! I'm your Tool-Calling Agent. I can help you with:\n\n"
                    f"🔢 **Calculator** — evaluate math expressions\n"
                    f"🌐 **Web Search** — look up information\n"
                    f"📁 **File Reader** — read files\n\n"
                    f"Try asking me something like:\n"
                    f"• \"What is 25 * 4 + 10?\"\n"
                    f"• \"Tell me about LangGraph\"\n"
                    f"• \"Read /kaggle/working/somefile.txt\"",
            duration_ms=(time.time() - t_start) * 1000,
        ))
        return steps
    
    tool_name = intent["tool"]
    tool_input = intent["input"]
    tool_info = TOOLS[tool_name]
    
    # Step 2: Thought about which tool
    steps.append(AgentStep(
        type=StepType.THOUGHT,
        content=f"I need to answer: '{query}'\n\n"
                f"→ Detected intent requires **{tool_name}**\n"
                f"→ Tool description: {tool_info['desc']}\n"
                f"→ Input: {tool_input[:100]}",
        duration_ms=thought_duration,
    ))
    
    # Step 3: Tool call
    steps.append(AgentStep(
        type=StepType.TOOL_CALL,
        content=f"Calling **{tool_name}** with input: `{tool_input[:100]}`",
        tool_name=tool_name,
        tool_input=tool_input[:100],
    ))
    
    # Step 4: Execute tool
    tool_start = time.time()
    try:
        result = tool_info["fn"](tool_input)
    except Exception as e:
        result = f"❌ Error executing {tool_name}: {e}"
    tool_duration = (time.time() - tool_start) * 1000
    
    steps.append(AgentStep(
        type=StepType.TOOL_RESULT,
        content=f"**{tool_name}** completed in {tool_duration:.0f}ms",
        tool_name=tool_name,
        tool_output=str(result)[:500],
        duration_ms=tool_duration,
    ))
    
    # Step 5: Final answer
    is_error = result.startswith("❌")
    if is_error:
        final = f"⚠️ The **{tool_name}** tool encountered an issue:\n\n```\n{result}\n```\n\n**Recovery:** Double-check your input and try again."
    else:
        final = f"Here's what I found using **{tool_name}**:\n\n{result}"
    
    steps.append(AgentStep(
        type=StepType.FINAL,
        content=final,
        duration_ms=(time.time() - t_start) * 1000,
    ))
    
    return steps


def format_steps_as_html(steps: List[AgentStep]) -> str:
    """Format agent steps as nice HTML for Gradio."""
    html = ""
    colors = {
        StepType.THOUGHT: ("#f0f4f8", "#2c3e50", "🧠"),
        StepType.TOOL_CALL: ("#fff3e0", "#e65100", "🔧"),
        StepType.TOOL_RESULT: ("#e8f5e9", "#1b5e20", "📥"),
        StepType.FINAL: ("#e3f2fd", "#0d47a1", "💬"),
    }
    
    for i, step in enumerate(steps):
        bg, color, icon = colors[step.type]
        label = step.type.value.replace("_", " ").title()
        
        html += f"""
        <div style="background:{bg}; border-left:4px solid {color}; border-radius:8px; 
                    padding:14px 18px; margin:10px 0; font-family:'Segoe UI',system-ui,sans-serif;">
            <div style="display:flex; align-items:center; gap:8px; margin-bottom:6px;">
                <span style="font-size:18px;">{icon}</span>
                <strong style="color:{color}; font-size:14px;">Step {i+1}: {label}</strong>
                {f'<span style="margin-left:auto; color:#999; font-size:12px;">{step.duration_ms:.0f}ms</span>' if step.duration_ms > 0 else ''}
            </div>
            <div style="color:#333; font-size:14px; line-height:1.6;">
                {step.content}
            </div>
        """
        
        if step.tool_name and step.tool_input:
            html += f"""
            <div style="background:rgba(0,0,0,0.04); border-radius:4px; padding:8px 12px; margin-top:8px; font-family:monospace; font-size:13px;">
                <span style="color:#666;">Tool:</span> {step.tool_name} | 
                <span style="color:#666;">Input:</span> {step.tool_input}
            </div>
            """
        
        if step.tool_output:
            html += f"""
            <div style="background:#1e1e1e; color:#d4d4d4; border-radius:4px; padding:10px 14px; margin-top:8px; font-family:monospace; font-size:13px; white-space:pre-wrap; max-height:200px; overflow-y:auto;">
                {step.tool_output[:500]}
            </div>
            """
        
        html += "</div>"
    
    return html


# ============================================================
# Gradio UI
# ============================================================

def respond(message: str, history: list):
    """Process a user message and return the response."""
    if not message.strip():
        return "", history
    
    steps = run_agent(message)
    html_output = format_steps_as_html(steps)
    
    # Add to history (type="messages" format)
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": html_output})
    return "", history


CSS = """
.gradio-container { max-width: 900px !important; margin: auto !important; }
.chatbot .user { background: #e3f2fd !important; }
.chatbot .assistant { background: transparent !important; }
footer { display: none !important; }
"""

with gr.Blocks(
    css=CSS,
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="indigo",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
    ),
    title="Tool-Calling AI Agent",
) as demo:
    gr.Markdown(
        """
        # 🛠️ Tool-Calling AI Agent
        
        **Built with LangGraph architecture** — watch the agent think, call tools, and respond.
        
        The agent follows a structured pipeline: **Thought → Tool Call → Tool Result → Final Answer**.
        
        ### Available Tools:
        | Tool | Description | Example |
        |------|-------------|--------|
        | 🔢 Calculator | Safe math evaluation | `25 * 4 + 10` |
        | 🌐 Web Search | Knowledge base lookup | `Tell me about LangGraph` |
        | 📁 File Reader | Read text files | `Read README.md` |
        """,
    )
    
    chatbot = gr.Chatbot(
        label="Agent Conversation",
        height=600,
        show_label=False,
        bubble_full_width=False,
        avatar_images=(None, "🧠"),
        type="messages",
    )
    
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask me anything... (e.g., 'What is 2+2?', 'Tell me about LangGraph', 'Read README.md')",
            show_label=False,
            container=False,
            scale=8,
        )
        send = gr.Button("Send", variant="primary", scale=1)
        clear = gr.ClearButton([msg, chatbot], scale=1)
    
    examples = gr.Examples(
        examples=[
            ["What is 2 + 2?"],
            ["Calculate the area of a circle with radius 7"],
            ["Tell me about LangGraph"],
            ["What is LangChain?"],
            ["Read app.py"],
            ["Calculate sqrt(144) + 50 * 3"],
            ["What is Hugging Face?"],
        ],
        inputs=[msg],
        label="Try these examples",
    )
    
    # Bind events
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    send.click(respond, [msg, chatbot], [msg, chatbot])

if __name__ == "__main__":
    demo.launch()