import os
import requests
import openai
import gradio as gr
from swarm import Swarm, Agent

# Fetch API keys from environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PROXYCURL_API_KEY = os.getenv("PROXYCURL_API_KEY")
FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY")

# Function to fetch LinkedIn data using Proxycurl API
def fetch_linkedin_data(linkedin_url):
    headers = {'Authorization': f'Bearer {PROXYCURL_API_KEY}'}
    response = requests.get(f"https://api.proxycurl.com/v1/linkedin/{linkedin_url}", headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": "Unable to fetch LinkedIn data"}

# Function to fetch company information using Firecrawl API
def fetch_company_info(company_url):
    headers = {'Authorization': f'Bearer {FIRECRAWL_API_KEY}'}
    response = requests.get(f"https://api.firecrawl.com/v1/scrape?url={company_url}", headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": "Unable to fetch company information"}

# Function to structure the email using the "Start with Why" model
def structure_email(user_data, linkedin_info, company_info):
    why = f"I am passionate about {company_info.get('mission', 'your mission')} because it aligns with my experience as {linkedin_info.get('current_role', 'a professional')}."
    how = f"My skills in {user_data['role']} match the requirements and goals of your organization."
    what = f"I can bring my experience in {linkedin_info.get('skills', 'relevant skills')} to help achieve {company_info.get('goal', 'your company goals')}."
    structured_input = f"{why}\n\n{how}\n\n{what}"
    return structured_input

# Function to generate email content using Nvidia Nemotron LLM
def generate_email_content(api_key, prompt):
    openai.api_key = api_key
    response = openai.Completion.create(
        model="nemotron-70b",
        prompt=prompt,
        max_tokens=500
    )
    return response.choices[0].text

# Function to validate the generated email for professional tone and completeness
def validate_email(email_content):
    return "Why" in email_content and "How" in email_content and "What" in email_content

# Define Agent B - The agent responsible for processing and refining the email
def transfer_to_email_generation():
    return email_agent_b

# Define Agent A - The primary agent to initiate data collection and hand over to Agent B
email_agent_a = Agent(
    name="Data Collection Agent",
    instructions="Collect user inputs and relevant data (LinkedIn and company details).",
    functions=[transfer_to_email_generation]
)

# Define Agent B - The agent that structures and generates the email
email_agent_b = Agent(
    name="Email Generation Agent",
    instructions="Structure the email using the 'Start with Why' model and generate professional content.",
)

# Set up the Swarm client
client = Swarm()

# Main function that interacts with the agents using the Swarm framework
def run_agent(name, email, phone, linkedin_url, company_url, role):
    user_data = {
        "name": name,
        "email": email,
        "phone": phone,
        "linkedin_url": linkedin_url,
        "company_url": company_url,
        "role": role
    }

    # Step 1: Fetch LinkedIn and Company Data
    linkedin_info = fetch_linkedin_data(linkedin_url)
    company_info = fetch_company_info(company_url)
    if "error" in linkedin_info or "error" in company_info:
        return "Error fetching data. Please check the LinkedIn and company URLs."

    # Step 2: Use the Swarm framework to structure and generate the email
    messages = [{"role": "user", "content": "Initiate email generation."}]
    response = client.run(agent=email_agent_a, messages=messages)

    # Step 3: Agent B structures and refines the email using the provided data
    prompt = structure_email(user_data, linkedin_info, company_info)
    email_content = generate_email_content(OPENAI_API_KEY, prompt)

    # Validate the generated email content using the ReAct pattern (with a max of 3 iterations)
    for i in range(3):
        if validate_email(email_content):
            return email_content
        else:
            # Refine the prompt based on the feedback or iteration logic
            refined_prompt = f"Refine: {prompt}"
            email_content = generate_email_content(OPENAI_API_KEY, refined_prompt)

    return "Unable to generate a valid email after 3 attempts."

# Set up the Gradio interface
final_interface = gr.Interface(
    fn=run_agent,
    inputs=[
        gr.Textbox(label="Name"),
        gr.Textbox(label="Email"),
        gr.Textbox(label="Phone Number"),
        gr.Textbox(label="LinkedIn Profile URL"),
        gr.Textbox(label="Company URL or Name"),
        gr.Textbox(label="Role Being Applied For")
    ],
    outputs="text",
    title="Email Writing AI Agent",
    description="Autonomously generate a professional email tailored to the job application."
)

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