import openai
import os

# Load API key from environment variable
openai.api_key = os.getenv("OPENAI_API_KEY")

def generate_classification_data(prompt, n_samples=10):
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt=f"Generate {n_samples} examples of mental health text and their categories. Categories can be depression, anxiety, stress, or happiness. Format: 'text -> category'.",
        max_tokens=150
    )
    
    generated_text = response.choices[0].text.strip().split("\n")
    data = [line.split(" -> ") for line in generated_text if "->" in line]
    
    return data

def classify_text(text):
    training_data = generate_classification_data("mental health text classification")
    
    X = [item[0] for item in training_data]
    y = [item[1] for item in training_data]
    
    if "sad" in text or "hopeless" in text:
        return "depression"
    elif "worried" in text or "anxious" in text:
        return "anxiety"
    elif "stressed" in text or "overwhelmed" in text:
        return "stress"
    else:
        return "happiness"