import streamlit as st
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from langchain_huggingface import HuggingFacePipeline
from langchain.agents import create_react_agent, AgentExecutor, Tool
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from langchain.schema import AgentAction, AgentFinish
import concurrent.futures
import time

# Email configuration
SENDER_EMAIL = "clen.emenike@gmail.com"
SENDER_PASSWORD = "Achuta@86"
SMTP_SERVER = "smtp.gmail.com"
SMTP_PORT = 587

@st.cache_resource
def load_model():
    model_name = "google/flan-t5-small"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_length=512)
    return HuggingFacePipeline(pipeline=pipe)

local_llm = load_model()

def send_email(to_email, subject, body):
    try:
        message = MIMEMultipart()
        message["From"] = SENDER_EMAIL
        message["To"] = to_email
        message["Subject"] = subject
        message.attach(MIMEText(body, "plain"))
        with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as server:
            server.starttls()
            server.login(SENDER_EMAIL, SENDER_PASSWORD)
            server.send_message(message)
        return "Email sent successfully"
    except Exception as e:
        return f"Failed to send email: {str(e)}"

# Define tools before creating the agent
tools = [Tool(name="Send Email", func=send_email, description="Sends an email. Args: to_email, subject, body")]

# Define the prompt before creating the agent
prompt = PromptTemplate.from_template(
    """You are a helpful assistant scheduling cybersecurity program meetings. 
    Collect the user's name, email, and preferred meeting date, then send a meeting invitation.
    Begin by asking for the user's name if you don't have it.
    Zoom link: https://us04web.zoom.us/j/73793374638?pwd=S0TEJ30da7dhQ8viOdafMzPfCVzoLJ.1
    Meeting ID: 737 9337 4638

    {chat_history}
    Human: {input}
    Assistant: Let's approach this step-by-step:
    {agent_scratchpad}

    Tools available:
    {tools}

    Tool names:
    {tool_names}
    """
)

# Create the agent after loading the model and defining tools and prompt
agent = create_react_agent(local_llm, tools, prompt)

# Create the AgentExecutor
agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent, 
    tools=tools, 
    verbose=True, 
    memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True),
    handle_parsing_errors=True,
    max_iterations=3,  # Limit the number of iterations
    early_stopping_method="force"  # or "generate_until_stop"
)

def run_agent_with_timeout(executor, input_data, timeout):
    with concurrent.futures.ThreadPoolExecutor() as pool:
        future = pool.submit(executor, input_data)
        try:
            return future.result(timeout=timeout)
        except concurrent.futures.TimeoutError:
            raise TimeoutError("The operation timed out.")

st.title("CyberSecurity Program Meeting Scheduler")
st.write("Chat with the AI to schedule your meeting. The AI will ask for your name, email, and preferred meeting date.")

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("Your message"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                start_time = time.time()
                response = run_agent_with_timeout(agent_executor, {"input": prompt}, timeout=10)
                if isinstance(response, dict) and "output" in response:
                    assistant_response = response["output"]
                elif isinstance(response, (AgentAction, AgentFinish)):
                    assistant_response = response.return_values.get("output", str(response))
                else:
                    assistant_response = str(response)
                
                st.markdown(assistant_response)
                st.session_state.messages.append({"role": "assistant", "content": assistant_response})
            except TimeoutError:
                st.error("I apologize, but I'm having trouble processing your request at the moment. Could you please try asking your question again, or rephrase it?")
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")

st.sidebar.title("About")
st.sidebar.info("This is an interactive CyberSecurity Program Meeting Scheduler. Chat with the AI to schedule your meeting. It will collect your information and send a meeting invitation via email.")

# To run this script, use: streamlit run your_script_name.py