from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    InferenceClientModel,
    load_tool,
    tool
)

import datetime
import pytz
import yaml

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# ---------------------------------------------------------
# CUSTOM TOOL 1: Current time
# ---------------------------------------------------------
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Gets the current local date and time for a timezone.

    Args:
        timezone: Valid timezone such as America/New_York,
                  Europe/London, or Asia/Kolkata.
    """

    try:
        tz = pytz.timezone(timezone)

        local_time = datetime.datetime.now(tz).strftime(
            "%Y-%m-%d %H:%M:%S"
        )

        return f"The current local time in {timezone} is {local_time}"

    except Exception as e:
        return f"Could not get time for {timezone}: {str(e)}"


# ---------------------------------------------------------
# CUSTOM TOOL 2: AI technology impact analyzer
# ---------------------------------------------------------
@tool
def analyze_ai_technology(topic: str) -> str:
    """
    Gives a quick engineering impact analysis for common AI technologies.

    Args:
        topic: AI technology or concept such as MCP, RAG,
               LangGraph, vector database, agents, or embeddings.
    """

    technologies = {
        "mcp": {
            "name": "Model Context Protocol (MCP)",
            "impact": 90,
            "description":
                "A standard protocol for connecting AI applications "
                "to tools, APIs, services, and external data.",
            "use_case":
                "Let an AI agent access GitHub, databases, files, "
                "Slack, APIs, and enterprise systems."
        },

        "rag": {
            "name": "Retrieval-Augmented Generation (RAG)",
            "impact": 88,
            "description":
                "Retrieves relevant information before asking an LLM "
                "to generate an answer.",
            "use_case":
                "Enterprise search, documentation assistants, "
                "knowledge bots, and support systems."
        },

        "langgraph": {
            "name": "LangGraph",
            "impact": 82,
            "description":
                "Framework for building stateful and multi-step "
                "LLM workflows and agents.",
            "use_case":
                "Agent orchestration, approval workflows, "
                "multi-agent systems, and long-running AI tasks."
        },

        "vector database": {
            "name": "Vector Database",
            "impact": 80,
            "description":
                "Stores embeddings and performs semantic similarity search.",
            "use_case":
                "RAG, recommendation systems, semantic search, "
                "and AI memory."
        },

        "embeddings": {
            "name": "Embeddings",
            "impact": 78,
            "description":
                "Numeric representations of text, images, or other data "
                "that capture semantic meaning.",
            "use_case":
                "Similarity search, clustering, recommendations, and RAG."
        },

        "agents": {
            "name": "AI Agents",
            "impact": 92,
            "description":
                "LLM-powered systems that can reason about a task "
                "and decide which tools or actions to execute.",
            "use_case":
                "Coding assistants, research agents, workflow automation, "
                "customer support, and enterprise automation."
        }
    }

    key = topic.lower().strip()

    # Handle a few common variations
    aliases = {
        "agent": "agents",
        "ai agent": "agents",
        "ai agents": "agents",
        "model context protocol": "mcp",
        "retrieval augmented generation": "rag",
        "vector db": "vector database",
        "vectors": "vector database",
        "embedding": "embeddings"
    }

    key = aliases.get(key, key)

    if key not in technologies:
        return (
            f"I don't have a predefined analysis for '{topic}'. "
            "Use web search to research it instead."
        )

    item = technologies[key]

    return f"""
Technology: {item['name']}
Estimated engineering impact: {item['impact']}%

What it is:
{item['description']}

Typical use case:
{item['use_case']}
"""


# ---------------------------------------------------------
# CUSTOM TOOL 3: Developer recommendation
# ---------------------------------------------------------
@tool
def recommend_ai_learning_topic(current_skill: str) -> str:
    """
    Suggests an AI engineering topic to learn based on a developer's skill.

    Args:
        current_skill: Developer skill such as React, Node.js,
                       Python, backend, cloud, or fullstack.
    """

    skill = current_skill.lower()

    if "react" in skill or "frontend" in skill:
        return (
            "Recommended next topic: AI application development.\n"
            "Learn LLM APIs, streaming responses, tool calling, "
            "AI SDKs, and agent UIs."
        )

    if "node" in skill or "backend" in skill:
        return (
            "Recommended next topic: Agent orchestration.\n"
            "Learn tool calling, MCP, RAG, LangGraph, queues, "
            "and asynchronous AI workflows."
        )

    if "python" in skill:
        return (
            "Recommended next topic: AI agent frameworks.\n"
            "Explore smolagents, LangGraph, PydanticAI, "
            "RAG pipelines, and evaluation."
        )

    if "cloud" in skill or "aws" in skill:
        return (
            "Recommended next topic: AI platform engineering.\n"
            "Learn model gateways, vector databases, observability, "
            "guardrails, GPU inference, and MCP servers."
        )

    if "fullstack" in skill:
        return (
            "Recommended path:\n"
            "1. LLM APIs\n"
            "2. Structured outputs\n"
            "3. Tool calling\n"
            "4. RAG\n"
            "5. MCP\n"
            "6. LangGraph\n"
            "7. Agent evaluation and observability"
        )

    return (
        "Start with LLM APIs, prompt engineering, structured outputs, "
        "tool calling, RAG, and then AI agents."
    )


# ---------------------------------------------------------
# FINAL ANSWER TOOL
# ---------------------------------------------------------
final_answer = FinalAnswerTool()


# ---------------------------------------------------------
# MODEL
# ---------------------------------------------------------
model = InferenceClientModel(
    max_tokens=2096,
    temperature=0.5,
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    custom_role_conversions=None,
)


# ---------------------------------------------------------
# TOOL FROM HUGGING FACE HUB
# ---------------------------------------------------------
image_generation_tool = load_tool(
    "agents-course/text-to-image",
    trust_remote_code=True
)


# ---------------------------------------------------------
# LOAD SYSTEM PROMPTS
# ---------------------------------------------------------
with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)


# ---------------------------------------------------------
# CREATE AGENT
# ---------------------------------------------------------
agent = CodeAgent(
    model=model,

    tools=[
        # Search current information from the web
        DuckDuckGoSearchTool(),

        # Custom tools
        get_current_time_in_timezone,
        analyze_ai_technology,
        recommend_ai_learning_topic,

        # Hugging Face Hub tool
        image_generation_tool,

        # Required final-answer tool
        final_answer,
    ],

    max_steps=8,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="AI Developer Assistant",

    description=(
        "An AI engineering assistant that can search the web, "
        "analyze AI technologies, recommend learning topics, "
        "check world times, and generate images."
    ),

    prompt_templates=prompt_templates
)


# ---------------------------------------------------------
# START GRADIO UI
# ---------------------------------------------------------
GradioUI(agent).launch()