import streamlit as st
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool
from tools.final_answer import FinalAnswerTool
import yaml

# Set page title
st.title('Cleaning Planner')
st.divider()

# Display instructions
st.text('By using this app you commit to cleaning for the amount of time you have entered. Please immediately start to clean once you have received the cleaning plan.')
st.divider()

# Take user input for available cleaning time
st.markdown('How much time do you have to commit to cleaning as soon as you receive your cleaning plan?')
amount_of_time_to_clean = st.slider('Available Cleaning Time in minutes', 10, 120, 0)
st.divider()

# Take user input for tasks
st.markdown('Please describe each of the tasks you would like to complete. For each task give an estimate for how long it will take you to complete the task.')

# Initialize session state for tasks
if 'tasks' not in st.session_state:
    st.session_state.tasks = []

# Input fields for task description and time estimate
task_desc = st.text_input("Task Description")
task_time = st.number_input("Time Estimate (min)", min_value=0)

# Button to add task
if st.button("Add Task"):
    new_task = {'description': task_desc, 'time': task_time}
    st.session_state.tasks.append(new_task)

# Display the list of tasks
for idx, task in enumerate(st.session_state.tasks, start=1):
    st.write(f"Task {idx}: {task['description']} - {task['time']} min")

st.divider()
st.write('An AI model will be used to generate a cleaning plan for you. Once you have received your cleaning plan, please immediately start cleaning.')

if st.button("Generate Cleaning Plan"):
    # AI model code
    agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=HfApiModel())

    # Initialize session state for cleaning plan
    if 'cleaning_plan' not in st.session_state:
        st.session_state.cleaning_plan = []

    # Create and Launch Agent
    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 configuration
    model_id = 'https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.3,  # Lower temperature for more consistent responses
        model_id=model_id,  # 'Qwen/Qwen2.5-Coder-32B-Instruct',
        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=[
            DuckDuckGoSearchTool(),
            final_answer
        ],
        max_steps=6,
        verbosity_level=1,
        grammar=None,
        planning_interval=None,
        name=None,
        description=None,
        prompt_templates=prompt_templates
    )

    agent.run(f"Generate a cleaning plan based on this text:\n{st.session_state.tasks}\nTime available: {amount_of_time_to_clean} minutes")