from smolagents import CodeAgent, HfApiModel, load_tool, tool
import yaml
import datetime

from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Tool to calculate BMI and provide health feedback
@tool
def calculate_bmi(weight: float, height: float) -> str:
    """Calculates BMI and provides a health status.
    Args:
        weight: Weight in kilograms.
        height: Height in meters.
    """
    if height <= 0 or weight <= 0:
        return "Invalid input. Please enter positive values for weight and height."

    bmi = weight / (height ** 2)
    status = ""
    
    if bmi < 18.5:
        status = "underweight"
    elif 18.5 <= bmi < 24.9:
        status = "normal weight"
    elif 25 <= bmi < 29.9:
        status = "overweight"
    else:
        status = "obese"

    return f"Your BMI is {bmi:.2f}, which is considered {status}."

# Tool to analyze symptoms and suggest possible diseases
@tool
def analyze_symptoms(symptoms: str) -> str:
    """Analyzes symptoms and suggests possible diseases.
    Args:
        symptoms: A comma-separated list of symptoms.
    """
    symptom_disease_mapping = {
        "fever, cough": "Possible Flu or COVID-19",
        "headache, nausea": "Possible Migraine or Food Poisoning",
        "chest pain, shortness of breath": "Possible Heart Disease - Seek immediate medical help!",
        "fatigue, weight loss": "Possible Thyroid Issues or Diabetes",
        "stomach pain, diarrhea": "Possible Food Poisoning or Gastroenteritis"
    }

    for key, disease in symptom_disease_mapping.items():
        if all(symptom in symptoms.lower() for symptom in key.split(", ")):
            return f"Based on your symptoms ({symptoms}), you may have: {disease}. Consult a doctor for a proper diagnosis."

    return "I'm not sure. Please consult a healthcare professional."

# Load model
final_answer = FinalAnswerTool()
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct'
)

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

# Define Agent
agent = CodeAgent(
    model=model,
    tools=[final_answer, calculate_bmi, analyze_symptoms],  # Added tools here
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="HealthcareBot",
    description="A healthcare assistant that calculates BMI and suggests potential diseases based on symptoms.",
    prompt_templates=prompt_templates
)

# Launch Gradio UI
GradioUI(agent).launch()