import os
import re

from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.output_parsers import RegexParser

from dotenv import load_dotenv
load_dotenv()

# Set up the OpenAI API key
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')

# Initialize the language model
llm = ChatOpenAI(model="gpt-4o-mini")

# Function to display model outputs
def display_output(output):
    """Display the model's output in a formatted manner."""
    print("Model Output:")
    print("-" * 40)
    print(output)
    print("-" * 40)
    print()

constrained_prompt = PromptTemplate(
    input_variables=["product", "target_audience", "tone", "word_limit"],
    template="""Create a product description for {product} targeted at {target_audience}.
    Use a {tone} tone and keep it under {word_limit} words.
    The description should include:
    1. A catchy headline
    2. Three key features
    3. A call to action
    
    Product Description:
    """
)

# Generate the constrained output
input_variables = {
    "product": "smart water bottle",
    "target_audience": "health-conscious millennials",
    "tone": "casual and friendly",
    "word_limit": "75"
}

chain = constrained_prompt | llm
output = chain.invoke(input_variables).content
display_output(output)

job_posting_prompt = PromptTemplate(
    input_variables=["job_title", "company", "location", "experience"],
    template="""Create a job posting for a {job_title} position at {company} in {location}.
    The candidate should have {experience} years of experience.
    Follow these rules:
    1. Start with a brief company description (2 sentences)
    2. List 5 key responsibilities, each starting with an action verb
    3. List 5 required qualifications, each in a single sentence
    4. End with a standardized equal opportunity statement
    
    Format the output as follows:
    COMPANY: [Company Description]
    
    RESPONSIBILITIES:
    - [Responsibility 1]
    - [Responsibility 2]
    - [Responsibility 3]
    - [Responsibility 4]
    - [Responsibility 5]
    
    QUALIFICATIONS:
    - [Qualification 1]
    - [Qualification 2]
    - [Qualification 3]
    - [Qualification 4]
    - [Qualification 5]
    
    EEO: [Equal Opportunity Statement]
    """
)

# Generate the rule-based output
input_variables = {
    "job_title": "Senior Software Engineer",
    "company": "TechInnovate Solutions",
    "location": "San Francisco, CA",
    "experience": "5+"
}

chain = job_posting_prompt | llm
output = chain.invoke(input_variables).content
display_output(output)

# Define a regex parser for structured output
regex_parser = RegexParser(
    regex=r"COMPANY:\s*([\s\S]*?)\n\s*RESPONSIBILITIES:\s*([\s\S]*?)\n\s*QUALIFICATIONS:\s*([\s\S]*?)\n\s*EEO:\s*([\s\S]*)",
    output_keys=["company_description", "responsibilities", "qualifications", "eeo_statement"]
)
# This regex pattern captures the company description, responsibilities, qualifications, and EEO statement from the output text.

# Create a new prompt template that includes the parser instructions
parsed_job_posting_prompt = PromptTemplate(
    input_variables=["job_title", "company", "location", "experience"],
    template="""Create a job posting for a {job_title} position at {company} in {location}.
    The candidate should have {experience} years of experience.
    Follow these rules:
    1. Start with a brief company description (2 sentences)
    2. List 5 key responsibilities, each starting with an action verb
    3. List 5 required qualifications, each in a single sentence
    4. End with a standardized equal opportunity statement
    
    Format the output EXACTLY as follows:
    COMPANY: [Company Description]
    
    RESPONSIBILITIES:
    - [Responsibility 1]
    - [Responsibility 2]
    - [Responsibility 3]
    - [Responsibility 4]
    - [Responsibility 5]
    
    QUALIFICATIONS:
    - [Qualification 1]
    - [Qualification 2]
    - [Qualification 3]
    - [Qualification 4]
    - [Qualification 5]
    
    EEO: [Equal Opportunity Statement]
    """
)

def clean_output(output):
    for key, value in output.items():
        if isinstance(value, str):
            # Remove leading/trailing whitespace and normalize newlines
            output[key] = re.sub(r'\n\s*', '\n', value.strip())
    return output

# Generate the parsed output
chain = parsed_job_posting_prompt | llm
raw_output = chain.invoke(input_variables).content

# Parse and clean the output
parsed_output = regex_parser.parse(raw_output)
cleaned_output = clean_output(parsed_output)

# Display the parsed output
print("Parsed Output:")
for key, value in cleaned_output.items():
    print(f"{key.upper()}:")
    print(value)
    print()

review_prompt = PromptTemplate(
    input_variables=["product", "rating", "pros", "cons", "word_limit"],
    template="""Write a product review for {product} with the following constraints:
    1. The review should have a {rating}-star rating (out of 5)
    2. Include exactly {pros} pros and {cons} cons
    3. Use between 2 and 3 sentences for each pro and con
    4. The entire review should be under {word_limit} words
    5. End with a one-sentence recommendation
    
    Format the review as follows:
    Rating: [X] out of 5 stars
    
    Pros:
    1. [Pro 1]
    2. [Pro 2]
    ...
    
    Cons:
    1. [Con 1]
    2. [Con 2]
    ...
    
    Recommendation: [One-sentence recommendation]
    """
)

# Generate the constrained review
input_variables = {
    "product": "Smartphone X",
    "rating": "4",
    "pros": "3",
    "cons": "2",
    "word_limit": "200"
}

chain = review_prompt | llm
output = chain.invoke(input_variables).content
display_output(output)