from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
# @tool
# def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
#     #Keep this format for the description / args / args description but feel free to modify the tool
#     """A tool that does nothing yet 
#     Args:
#         arg1: the first argument
#         arg2: the second argument
#     """
#     return "What magic will you build ?"


def __init__(self):
    self.patient_data = {}

@tool
def collect_symptoms(patient_id: str, symptoms: str) -> str:
    """
    Nurse tool to collect symptoms and ask follow-up questions.

    Args:
        patient_id: Unique identifier for the patient.
        symptoms: Initial symptoms provided by the patient.

    Returns:
        Follow-up questions based on the symptoms.
    """
    try:
        if not patient_id or not symptoms:
            raise ValueError("Patient ID and symptoms must be provided.")

        self.patient_data[patient_id] = {"symptoms": symptoms, "additional_info": {}}

        questions = [
            "How long have you had these symptoms?",
            "Do you have any allergies?",
            "Are you taking any medications?",
            "Have you experienced these symptoms before?",
            "Have you had any recent illnesses?",
            "Have you noticed any other unusual changes?",
            "What is your medical history related to these symptoms?"
            
        ]

        return f"Nurse: I have noted the symptoms ({symptoms}). Here are follow-up questions:\n" + "\n".join(questions)

    except ValueError as e:
        return f"Error: {str(e)}"
    except Exception as e:
        return f"Unexpected Error: {str(e)}"

@tool
def diagnose_patient(patient_id: str) -> str:
    """
    Doctor tool to diagnose the patient based on symptoms.
    
    Args:
        patient_id: Unique identifier for the patient.
        
    Returns:
        Diagnosis and recommendations.
    """
    try:
        if patient_id not in CodeAgent.patient_data:
            raise ValueError("No symptoms found. Nurse must collect symptoms first.")

        symptoms = CodeAgent.patient_data[patient_id]["symptoms"]
        diagnosis = CodeAgent.fetch_diagnosis(symptoms)
        medication = CodeAgent.fetch_medication(symptoms)
        advice = CodeAgent.fetch_treatment_advice(symptoms)

        return (f"Doctor: Based on the symptoms: {symptoms},\n"
                f"Diagnosis: {diagnosis}\n"
                f"Medication: {medication}\n"
                f"Advice: {advice}")

    except ValueError as e:
        return f"Error: {str(e)}"
    except Exception as e:
        return f"Unexpected Error: {str(e)}"


@tool
def fetch_diagnosis(symptoms: str) -> str:
    """
    AI tool to retrieve a diagnosis based on symptoms.

    Args:
        symptoms: The symptoms provided by the patient.

    Returns:
        Detailed Diagnosis information.
    """
    search_query = f" A detailed medical diagnosis for this {symptoms}"
    return search_tool(search_query)
    try:
        if not symptoms:
            raise ValueError("Symptoms must be provided.")

        return f"AI Diagnosis: Potential cause of {symptoms} found."

    except ValueError as e:
        return f"Error: {str(e)}"
    except Exception as e:
        return f"Unexpected Error: {str(e)}"

@tool
def fetch_medication(symptoms: str) -> str:
    """
    AI tool to suggest medications based on symptoms.

    Args:
        symptoms: The symptoms provided by the patient.

    Returns:
        Suggested medication.
    """
    search_query = f" A detailed suggested medication for this {symptoms}"
    return search_tool(search_query)
    try:
        if not symptoms:
            raise ValueError("Symptoms must be provided.")

        return f"Suggested medication for {symptoms} found."

    except ValueError as e:
        return f"Error: {str(e)}"
    except Exception as e:
        return f"Unexpected Error: {str(e)}"

@tool
def fetch_treatment_advice(symptoms: str) -> str:
    """
    AI tool to provide treatment recommendations.

    Args:
        symptoms: The symptoms provided by the patient.

    Returns:
        Recommended treatment and advice.
    """
    search_query = f" A detailed recommended treatment and advice for this {symptoms}"
    return search_tool(search_query)
    try:
        if not symptoms:
            raise ValueError("Symptoms must be provided.")

        return f"Treatment advice for {symptoms} retrieved."

    except ValueError as e:
        return f"Error: {str(e)}"
    except Exception as e:
        return f"Unexpected Error: {str(e)}"
        
search_tool = DuckDuckGoSearchTool()

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that 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 the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='meta-llama/Llama-3.1-8B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, collect_symptoms, diagnose_patient, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="MedicalDoctor",
    description="A tool that simulates the roles of a medical doctor and a nurse in gathering symptoms, diagnosing, and providing recommendations",
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()