import os
import asyncio
import nest_asyncio
from dotenv import load_dotenv
import gradio as gr
from contextlib import contextmanager
from openai import OpenAI
import sendgrid
from sendgrid.helpers.mail import Mail, Email, To, Content
from pathlib import Path
import certifi

os.environ['SSL_CERT_FILE'] = certifi.where()

# Patch event loop for notebook compatibility
nest_asyncio.apply()

env_path = Path(__file__).resolve().parents[1] / ".env"
load_dotenv(dotenv_path=env_path, override=True)

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# --- Simple trace context manager for logging ---
@contextmanager
def trace(name):
    print(f"TRACE START: {name}")
    yield
    print(f"TRACE END: {name}")

# --- Agent class ---
class Agent:
    def __init__(self, name, instructions, model="gpt-4o-mini", tools=None, handoffs=None):
        self.name = name
        self.instructions = instructions
        self.model = model
        self.tools = tools or []
        self.handoffs = handoffs or []

# --- Runner logic ---
class Runner:
    @staticmethod
    async def run(agent, input_text):
        print(f"Running agent: {agent.name}")
        # The OpenAI client call is synchronous here (no await)
        response = client.chat.completions.create(
            model=agent.model,
            messages=[
                {"role": "system", "content": agent.instructions},
                {"role": "user", "content": input_text}
            ]
        )
        content = response.choices[0].message.content.strip()
        return type("Result", (), {"final_output": content})

# --- Sales agents instructions ---
instructions1 = (
    "You are a sales agent working for ComplAI, a company that provides a SaaS tool for ensuring SOC2 compliance "
    "and preparing for audits, powered by AI. You write professional, serious cold emails."
)
instructions2 = (
    "You are a humorous, engaging sales agent working for ComplAI, a company that provides a SaaS tool for "
    "ensuring SOC2 compliance and preparing for audits, powered by AI. You write attention-grabbing, lighthearted cold emails."
)
instructions3 = (
    "You are a busy sales agent who writes very short, direct cold emails for ComplAI, a SaaS compliance automation tool."
)

sales_agent1 = Agent("Professional Sales Agent", instructions1)
sales_agent2 = Agent("Engaging Sales Agent", instructions2)
sales_agent3 = Agent("Busy Sales Agent", instructions3)

sales_picker = Agent(
    name="Sales Picker",
    instructions=(
        "You pick the best cold sales email from the given options. "
        "Imagine you are a customer and pick the one you are most likely to respond to. "
        "Do not give an explanation; reply with the selected email only."
    )
)

# --- SendGrid email sending function ---
def send_email(body: str):
    sg = sendgrid.SendGridAPIClient(api_key=os.getenv('SENDGRID_API_KEY'))
    from_email = Email("jonathand2028@uchicago.edu")  # Change to your verified sender email
    to_email = To("jonathand2028@uchicago.edu")      # Change to your recipient email
    content = Content("text/plain", body)
    mail = Mail(from_email, to_email, "Sales email", content)
    response = sg.client.mail.send.post(request_body=mail.get())
    return {"status": "success", "code": response.status_code}

# --- Async functions to generate emails and pick best ---
async def generate_emails_async(product_description: str):
    with trace("Parallel cold emails"):
        results = await asyncio.gather(
            Runner.run(sales_agent1, product_description),
            Runner.run(sales_agent2, product_description),
            Runner.run(sales_agent3, product_description),
        )
    return [r.final_output for r in results]

async def pick_best_email_async(emails):
    email_block = "\n\n".join(f"Option {i+1}:\n{email}" for i, email in enumerate(emails))
    selection_result = await Runner.run(sales_picker, email_block)
    return selection_result.final_output

# --- Format emails nicely for Markdown display ---
def format_emails_markdown(emails, best_email):
    md = ""
    for i, email in enumerate(emails, 1):
        md += f"### Agent {i} Email:\n\n```\n{email.strip()}\n```\n\n"
    md += f"## 🟢 Best Picked Email:\n\n```\n{best_email.strip()}\n```\n"
    return md

# --- Sync wrapper for Gradio ---
def generate_and_select(product_description: str):
    loop = asyncio.get_event_loop()
    emails = loop.run_until_complete(generate_emails_async(product_description))
    best = loop.run_until_complete(pick_best_email_async(emails))
    send_result = send_email(best)
    formatted_md = format_emails_markdown(emails, best)
    return formatted_md, ""  # Clear input box

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 💼 AI Sales Email Generator with Email Sender")
    gr.Markdown(
        "Enter a product or service below. Three AI sales reps will write cold emails, "
        "a fourth AI will pick the best one, and it will be sent automatically."
    )

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Product or Service", placeholder="e.g. AI compliance tool for startups")
            generate_btn = gr.Button("Generate Emails & Send")
        with gr.Column():
            output_text = gr.Markdown(label="Results")

    generate_btn.click(
        fn=generate_and_select,
        inputs=input_text,
        outputs=[output_text, input_text]
    )

if __name__ == "__main__":
    demo.launch()