from smolagents import CodeAgent, HfApiModel, tool
import datetime
import pytz
import yaml
import gradio as gr
import os
from tools.final_answer import FinalAnswerTool

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

    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').

    Returns:
        A string with the current local time in the specified timezone, or an error message if the timezone is invalid.
    """
    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"Error fetching time for timezone '{timezone}': {str(e)}"

final_answer = FinalAnswerTool()

if not os.getenv("HF_TOKEN"):
    raise ValueError("HF_TOKEN environment variable not set. Please add it as a secret in Spaces settings.")

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

try:
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
        custom_role_conversions=None,
    )
except Exception as e:
    print(f"Error with Qwen model: {e}. Falling back to alternative endpoint.")
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
        custom_role_conversions=None,
    )

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone],
    max_steps=6,
    prompt_templates=prompt_templates
)

def chat_with_agent(message, history):
    try:
        print(f"Processing message: {message}")
        response = agent.run(message) if hasattr(agent, 'run') else str(agent)
        print(f"Raw response: {response}")
        if isinstance(response, dict):
            return response.get("output", response.get("final_answer", str(response)))
        return str(response)
    except Exception as e:
        print(f"Error in chat_with_agent: {e}")
        return f"Error: {str(e)}"

demo = gr.ChatInterface(
    fn=chat_with_agent,
    title="CodeAgent Chat",
    description="Interact with a problem-solving agent.",
)

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