import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import requests
import re
from datetime import datetime
from typing import List
import gradio as gr

# ============================================
# TOOLS
# ============================================

def web_search(query: str) -> str:
    """Search the web"""
    try:
        url = f"https://api.duckduckgo.com/?q={query}&format=json&no_html=1"
        r = requests.get(url, timeout=10)
        data = r.json()
        results = []
        for topic in data.get("RelatedTopics", [])[:3]:
            if isinstance(topic, dict) and "Text" in topic:
                results.append(topic["Text"][:300])
        return "\n".join(results) if results else "No results found"
    except Exception as e:
        return f"Search failed: {str(e)}"

def wikipedia(topic: str) -> str:
    """Get Wikipedia summary"""
    try:
        url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{topic.replace(' ', '_')}"
        r = requests.get(url, timeout=10)
        data = r.json()
        return data.get("extract", "Not found")[:1000]
    except Exception as e:
        return f"Wikipedia lookup failed: {str(e)}"

def get_weather(city: str) -> str:
    """Get weather for a city"""
    try:
        r = requests.get(f"https://wttr.in/{city}?format=%C:+%t", timeout=10)
        return f"Weather in {city}: {r.text.strip()}"
    except Exception as e:
        return f"Weather lookup failed: {str(e)}"

def calculate(expression: str) -> str:
    """Calculate math expression"""
    expression = re.sub(r'[^0-9+\-*/().%]', '', expression)
    if not expression:
        return "Please provide a math expression like '2+2'"
    try:
        result = eval(expression)
        return f"{expression} = {result}"
    except Exception as e:
        return f"Calculation error: {str(e)}"

def get_time() -> str:
    """Get current time"""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def summarize_text(text: str) -> str:
    """Simple summarization"""
    if len(text) < 50:
        return "Text too short to summarize."
    sentences = re.split(r'[.!?]+', text)
    key_sentences = [s.strip() for s in sentences if len(s.strip()) > 30][:3]
    if not key_sentences:
        return text[:300] + "..."
    return ". ".join(key_sentences) + "."

# Tool mapping
TOOLS = {
    "web_search": web_search,
    "wikipedia": wikipedia,
    "weather": get_weather,
    "calculate": calculate,
    "current_time": get_time,
    "summarize": summarize_text,
}

# ============================================
# ROUTER
# ============================================

def route_task(task: str) -> List[str]:
    """Simple keyword-based router"""
    task_lower = task.lower()
    routing = {
        "web_search": ["search", "google", "find", "look up", "internet", "news"],
        "wikipedia": ["wiki", "wikipedia", "who is", "what is"],
        "weather": ["weather", "temperature", "rain", "sunny"],
        "calculate": ["calculate", "math", "plus", "minus", "times", "divide"],
        "current_time": ["time", "date", "today", "now"],
        "summarize": ["summarize", "summary", "shorten"],
    }
    
    selected = []
    for tool, keywords in routing.items():
        if any(k in task_lower for k in keywords):
            selected.append(tool)
    
    return selected[:3] if selected else ["web_search"]

# ============================================
# LOAD MODEL
# ============================================

print("📥 Loading model...")
tokenizer = AutoTokenizer.from_pretrained("sirev/Gemma-2b-Uncensored-v1")
model = AutoModelForCausalLM.from_pretrained(
    "sirev/Gemma-2b-Uncensored-v1",
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True
)
print("✅ Model loaded!")

def generate_response(prompt: str) -> str:
    """Generate response from the model"""
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
    outputs = model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = response[len(prompt):].strip()
    return response if response else "I couldn't generate a response."

# ============================================
# AGENT FUNCTION
# ============================================

def run_agent(user_query: str) -> str:
    """Main agent function"""
    if not user_query or user_query.strip() == "":
        return "Please enter a question."
    
    print(f"Processing: {user_query}")
    
    # Route to find relevant tools
    tools_to_use = route_task(user_query)
    print(f"Selected tools: {tools_to_use}")
    
    # Execute tools
    tool_results = []
    for tool in tools_to_use:
        try:
            if tool == "web_search":
                result = web_search(user_query)
            elif tool == "wikipedia":
                words = user_query.split()
                topic = ' '.join(words[:4]) if words else user_query[:50]
                result = wikipedia(topic)
            elif tool == "weather":
                city = "London"
                if "in" in user_query.split():
                    words = user_query.split()
                    city = words[words.index("in") + 1]
                result = get_weather(city)
            elif tool == "calculate":
                math_match = re.search(r'[\d\s\+\-\*/\(\)\.\%]+', user_query)
                expr = math_match.group(0) if math_match else user_query
                result = calculate(expr)
            elif tool == "current_time":
                result = get_time()
            elif tool == "summarize":
                result = summarize_text(user_query)
            else:
                result = f"Tool '{tool}' not implemented"
            
            tool_results.append(f"🔹 {tool}: {result[:400]}")
            print(f"✓ {tool} executed")
        except Exception as e:
            tool_results.append(f"🔸 {tool}: Error - {str(e)[:100]}")
            print(f"✗ {tool} failed: {e}")
    
    # Generate response
    tool_context = "\n".join(tool_results) if tool_results else "No tools executed."
    
    context = f"""User asked: {user_query}

Tool results:
{tool_context}

Please answer based on the tool results above."""
    
    print("Generating response...")
    response = generate_response(context)
    
    # Add tool info prefix
    if tool_results:
        final = f"🔧 {', '.join(tools_to_use)}\n\n{response}"
    else:
        final = response
    
    print("Done")
    return final

# ============================================
# GRADIO INTERFACE - WORKING VERSION
# ============================================

examples = [
    "What is artificial intelligence?",
    "Search for AI news",
    "Weather in Tokyo",
    "Calculate 25 * 4",
    "What time is it?",
]

with gr.Blocks(title="Tool-Augmented AI Assistant") as demo:
    gr.Markdown("# 🛠️ Tool-Augmented AI Assistant")
    
    chatbot = gr.Chatbot(label="Conversation", height=500)
    
    with gr.Row():
        msg = gr.Textbox(label="Your Question", scale=4)
        submit = gr.Button("Send", variant="primary", scale=1)
    
    clear = gr.Button("Clear")
    gr.Examples(examples, inputs=msg)
    
    state = gr.State([])
    
    def respond(message, history):
        if not message:
            return "", history
        
        response = run_agent(message)
        history.append((message, response))
        return "", history
    
    def reset():
        return [], []
    
    msg.submit(respond, [msg, state], [msg, state])
    submit.click(respond, [msg, state], [msg, state])
    clear.click(reset, None, [chatbot, state])

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