import os

import yaml
from dotenv import load_dotenv
from smolagents import CodeAgent, InferenceClientModel, load_tool

from Gradio_UI import GradioUI
from tools.final_answer import FinalAnswerTool
from tools.get_current_time_in_timezone import get_current_time_in_timezone
from tools.get_weather import get_weather
from tools.search_wikipedia import search_wikipedia
from tools.visit_webpage import VisitWebpageTool
from tools.web_search import DuckDuckGoSearchTool

load_dotenv()

# --- Core Tools & Agent Setup ---

final_answer = FinalAnswerTool()
web_search = DuckDuckGoSearchTool()
visit_webpage = VisitWebpageTool()

hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    raise RuntimeError(
        "HF_TOKEN is required. Local model execution is disabled in this project."
    )

# Image generation tool from Hugging Face Hub
image_generation_tool = None
try:
    image_generation_tool = load_tool(
        "agents-course/text-to-image",
        trust_remote_code=True,
        token=hf_token,
    )
except Exception as e:
    print(f"Notice: text-to-image tool could not be loaded: {e}")

# Model configuration (cloud by default; Ollama supported for local learning/testing)
model_id = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct")
if "ollama" in model_id.lower():
    from smolagents import LiteLLMModel

    model = LiteLLMModel(
        model_id=model_id,
        api_base=os.getenv("OLLAMA_API_BASE", "http://127.0.0.1:11434"),
        num_ctx=int(os.getenv("OLLAMA_NUM_CTX", "8192")),
        temperature=0.5,
    )
else:
    model = InferenceClientModel(
        model_id=model_id,
        token=hf_token,
        max_tokens=2096,
        temperature=0.5,
        custom_role_conversions=None,
    )

# Load prompt templates
prompts_file = os.path.join(os.path.dirname(__file__), "prompts.yaml")
with open(prompts_file, "r") as stream:
    prompt_templates = yaml.safe_load(stream)

# --- Specialist Agents (Unit 2: Hierarchical Multi-Agent System) ---

# Specialist 1: Web & Wikipedia Researcher
research_agent = CodeAgent(
    model=model,
    tools=[web_search, search_wikipedia, visit_webpage],
    max_steps=5,
    verbosity_level=1,
    name="web_researcher",
    description="A specialist researcher that searches DuckDuckGo, queries Wikipedia, and reads webpage content to gather accurate facts and summaries. ALWAYS call this agent with a 'task' argument, e.g. web_researcher(task='Find the capital of Japan').",
    prompt_templates=prompt_templates
)

# Specialist 2: Weather & Timezone Specialist
local_info_agent = CodeAgent(
    model=model,
    tools=[get_current_time_in_timezone, get_weather],
    max_steps=4,
    verbosity_level=1,
    name="local_info_specialist",
    description="A specialist agent that fetches real-time live weather reports and calculates current time in timezones worldwide. ALWAYS call this agent with a 'task' argument, e.g. local_info_specialist(task='Get current weather in Tokyo'). Never call weather() or get_weather() directly from the manager.",
    prompt_templates=prompt_templates
)

# Make managed agents tolerant of common argument names (query/question/etc.)
def _adapt_managed_agent(agent):
    original_call = agent.__call__

    def _call(task=None, **kwargs):
        if task is None:
            for key in ("task", "query", "question", "request", "prompt"):
                if key in kwargs:
                    task = kwargs.pop(key)
                    break
            if task is None and kwargs:
                task = list(kwargs.values())[0]
                kwargs = {}
        return original_call(task, **kwargs)

    agent.__call__ = _call
    return agent

research_agent = _adapt_managed_agent(research_agent)
local_info_agent = _adapt_managed_agent(local_info_agent)

# Manager Agent: Coordinates specialists and synthesizes the final response
manager_tools = [final_answer]
if image_generation_tool is not None:
    manager_tools.append(image_generation_tool)

manager_agent = CodeAgent(
    model=model,
    tools=manager_tools,
    managed_agents=[research_agent, local_info_agent],
    max_steps=8,
    verbosity_level=1,
    planning_interval=None,
    name="ArkAssistant_Manager",
    description="A senior manager agent that breaks down user tasks and delegates them to specialized team members (web researcher and local info specialist).",
    prompt_templates=prompt_templates
)

# Main agent interface for Gradio
agent = manager_agent

if __name__ == "__main__":
    GradioUI(agent).launch()