# -*- coding: utf-8 -*-
"""
Extended smolagents template
Adds a tool that finds the overlap in normal office-hours
(09:00-17:00 local time by default) for a list of time-zones,
so a distributed team can quickly see when they’re all online.

Teams: Kyrgyzstan (Asia/Bishkek), USA (pick any valid TZ,
e.g. America/New_York), Uzbekistan (Asia/Tashkent).

Usage inside the chat UI, for example:
    find_overlapping_work_hours(
        ["Asia/Bishkek", "America/New_York", "Asia/Tashkent"],
        start_local="09:00",
        end_local="17:00"
    )
"""

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    InferenceClientModel,
    load_tool,
    tool,
)
import datetime
import pytz
import yaml
from typing import List

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# --------------------------------------------------------------------------- #
# Example placeholder tool (left intact)
def my_custom_tool(arg1: str, arg2: int) -> str:
    """A tool that does nothing yet
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"


# --------------------------------------------------------------------------- #
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Return the current wall-clock time for a given timezone.
    Args:
        timezone: IANA tz database string, e.g. 'America/New_York'.
    """
    try:
        tz = pytz.timezone(timezone)
        now_local = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Local time in {timezone}: {now_local}"
    except Exception as exc:
        return f"Error: {exc}"


# --------------------------------------------------------------------------- #
@tool
def find_overlapping_work_hours(
    timezones: List[str],
    start_local: str = "09:00",
    end_local: str = "17:00",
) -> str:
    """Given several IANA time-zones, return the daily overlap of office hours.

    Args:
        timezones: List of tz names (e.g. ['Asia/Bishkek','America/New_York'])
        start_local: Start of work day in HH:MM (24 h) for *each* zone
        end_local:   End   of work day in HH:MM (24 h) for *each* zone
    """
    try:
        # Parse the local start/end once
        start_h, start_m = map(int, start_local.split(":"))
        end_h, end_m = map(int, end_local.split(":"))
        if (end_h, end_m) <= (start_h, start_m):
            return "End time must be after start time."

        # For today’s date we’ll convert each zone’s window to UTC
        today = datetime.date.today()
        utc = pytz.utc
        earliest_end_utc = datetime.datetime.min.replace(tzinfo=utc)
        latest_start_utc = datetime.datetime.max.replace(tzinfo=utc)

        details = []
        for tz_name in timezones:
            tz = pytz.timezone(tz_name)

            local_start = tz.localize(
                datetime.datetime(today.year, today.month, today.day, start_h, start_m)
            )
            local_end = tz.localize(
                datetime.datetime(today.year, today.month, today.day, end_h, end_m)
            )

            start_utc = local_start.astimezone(utc)
            end_utc = local_end.astimezone(utc)

            # track overlap
            if start_utc > latest_start_utc:
                latest_start_utc = start_utc
            if end_utc < earliest_end_utc or earliest_end_utc == datetime.datetime.min.replace(
                tzinfo=utc
            ):
                earliest_end_utc = end_utc

            details.append(
                f"{tz_name}: {local_start.strftime('%H:%M')}–{local_end.strftime('%H:%M')} "
                f"(UTC {start_utc.strftime('%H:%M')}–{end_utc.strftime('%H:%M')})"
            )

        if earliest_end_utc <= latest_start_utc:
            overlap_msg = "No common working window today."
        else:
            # Present the intersection in UTC and in each local zone for clarity
            overlap_local = []
            for tz_name in timezones:
                tz = pytz.timezone(tz_name)
                overlap_start_local = latest_start_utc.astimezone(tz).strftime("%H:%M")
                overlap_end_local = earliest_end_utc.astimezone(tz).strftime("%H:%M")
                overlap_local.append(f"{tz_name}: {overlap_start_local}–{overlap_end_local}")

            overlap_msg = (
                f"✅ Overlap (UTC): {latest_start_utc.strftime('%H:%M')}–"
                f"{earliest_end_utc.strftime('%H:%M')}\n"
                + "\n".join(overlap_local)
            )

        return (
            "Daily office hours\n"
            + "\n".join(details)
            + "\n\n"
            + overlap_msg
        )
    except Exception as exc:
        return f"Error computing overlap: {exc}"


# --------------------------------------------------------------------------- #
# Required final-answer tool
final_answer = FinalAnswerTool()

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

# Optional extra tool from the Hugging Face Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# System prompt templates
with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

# Assemble the agent
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        get_current_time_in_timezone,
        find_overlapping_work_hours,
        # my_custom_tool,            # uncomment if you actually need it
        # image_generation_tool,     # idem
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
)

# Launch a small Gradio front-end
GradioUI(agent).launch()