import datetime
from pathlib import Path

import pytz
import yaml

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

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


# ============================================================
# 1. CUSTOM TOOL: HVAC AIRFLOW UNIT CONVERSION
# ============================================================

@tool
def convert_airflow_units(airflow_m3h: float) -> str:
    """
    Convert HVAC airflow from cubic metres per hour into
    litres per second and cubic feet per minute.

    Args:
        airflow_m3h: Airflow in cubic metres per hour.
    """
    if airflow_m3h < 0:
        return "Error: Airflow cannot be negative."

    airflow_ls = airflow_m3h / 3.6
    airflow_cfm = airflow_m3h * 0.588578

    return (
        f"Airflow: {airflow_m3h:.2f} m3/h\n"
        f"Airflow: {airflow_ls:.2f} L/s\n"
        f"Airflow: {airflow_cfm:.2f} CFM"
    )


# ============================================================
# 2. CUSTOM TOOL: CURRENT TIME IN A TIME ZONE
# ============================================================

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Get the current local date and time in a specified timezone.

    Args:
        timezone: Valid timezone, e.g. Asia/Riyadh or America/New_York.
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz)

        return (
            f"Current time in {timezone}: "
            f"{local_time.strftime('%Y-%m-%d %H:%M:%S')}"
        )

    except pytz.exceptions.UnknownTimeZoneError:
        return f"Error: Unknown timezone '{timezone}'."


# ============================================================
# 3. LOAD FINAL ANSWER TOOL
# ============================================================

final_answer = FinalAnswerTool()


# ============================================================
# 4. LOAD WEB SEARCH TOOL
# ============================================================

search_tool = DuckDuckGoSearchTool(max_results=5)


# ============================================================
# 5. LOAD IMAGE GENERATION TOOL
# ============================================================

# This tool comes from the Hugging Face Hub.
# Only load it if you trust the repository's code.

image_generation_tool = load_tool(
    "agents-course/text-to-image",
    trust_remote_code=True,
)


# ============================================================
# 6. LOAD PROMPT TEMPLATES
# ============================================================

BASE_DIR = Path(__file__).resolve().parent
PROMPTS_FILE = BASE_DIR / "prompts.yaml"

with PROMPTS_FILE.open("r", encoding="utf-8") as stream:
    prompt_templates = yaml.safe_load(stream)

if not isinstance(prompt_templates, dict):
    raise ValueError("prompts.yaml must contain a YAML dictionary.")


# ============================================================
# 7. INITIALIZE HUGGING FACE MODEL
# ============================================================

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
)


# ============================================================
# 8. CREATE THE AI AGENT
# ============================================================

agent = CodeAgent(
    model=model,

    tools=[
        final_answer,
        search_tool,
        convert_airflow_units,
        get_current_time_in_timezone,
        image_generation_tool,
    ],

    max_steps=8,
    verbosity_level=1,

    planning_interval=None,
    name="pdf_engineer_assistant",

    description=(
        "An AI assistant that searches the web, "
        "converts HVAC airflow units, checks time zones, "
        "and generates images."
    ),

    prompt_templates=prompt_templates,
)


# ============================================================
# 9. LAUNCH GRADIO INTERFACE
# ============================================================

if __name__ == "__main__":
    GradioUI(agent).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )