import os
import torch
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from langchain import PromptTemplate, LLMChain
from langchain.llms import HuggingFacePipeline
from langchain.agents import AgentExecutor, Tool, ZeroShotAgent
from langchain.memory import ConversationBufferMemory
import pandas as pd
from sqlalchemy import create_engine
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

# Set up the open-source LLM
@st.cache_resource
def load_model():
    model_name = "google/flan-t5-large"
    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()

# Set up the database connection
db_connection_string = "sqlite:///leads.db"  # Replace with your actual database connection string
engine = create_engine(db_connection_string)

# Define the tools for the agent
def search_leads(query):
    df = pd.read_sql(f"SELECT * FROM leads WHERE name LIKE '%{query}%'", engine)
    return df.to_dict(orient='records')

def send_email(to_email, subject, body):
    # For demo purposes, we'll just print the email details
    st.write(f"Email sent to: {to_email}")
    st.write(f"Subject: {subject}")
    st.write(f"Body: {body}")
    return "Email sent successfully"

tools = [
    Tool(
        name="Search Leads",
        func=search_leads,
        description="Useful for searching leads in the database"
    ),
    Tool(
        name="Send Email",
        func=send_email,
        description="Useful for sending emails to leads"
    )
]

# Set up the agent
prefix = """You are an AI CyberSecurity Program Advisor. Your goal is to engage with leads and get them to book a video call for an in-person sales meeting. You have access to a database of leads and can send emails.

You have access to the following tools:"""

suffix = """Begin!

{chat_history}
Human: {human_input}
AI: Let's approach this step-by-step:"""

prompt = ZeroShotAgent.create_prompt(
    tools,
    prefix=prefix,
    suffix=suffix,
    input_variables=["human_input", "chat_history"]
)

llm_chain = LLMChain(llm=local_llm, prompt=prompt)
memory = ConversationBufferMemory(memory_key="chat_history")

agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent, tools=tools, verbose=True, memory=memory
)

# Streamlit interface
st.title("AI CyberSecurity Program Advisor Demo")

st.write("This demo showcases an AI agent that can engage with leads and attempt to book video calls for in-person sales meetings.")

lead_name = st.text_input("Enter a lead's name to engage with:")

if lead_name:
    lead_info = search_leads(lead_name)
    if not lead_info:
        st.write(f"No lead found with the name {lead_name}")
    else:
        lead = lead_info[0]
        st.write(f"Lead found: {lead['name']} (Email: {lead['email']})")
        
        initial_message = f"Hello {lead['name']}, I'd like to discuss our cybersecurity program with you. Are you available for a quick video call?"
        
        if st.button("Engage with Lead"):
            with st.spinner("AI is generating a response..."):
                response = agent_executor.run(initial_message)
            
            st.write("AI Response:")
            st.write(response)

st.sidebar.title("About")
st.sidebar.info("This is a demo of an AI CyberSecurity Program Advisor using an open-source LLM and LangChain. It's designed to engage with leads and attempt to book video calls for sales meetings.")

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