import json
from groq import Groq
import os

class VariableExtractor:
    def __init__(self):
        self.client = None
    
    def set_api_key(self, api_key):
        """Set Groq API key"""
        self.client = Groq(api_key=api_key)
    
    def extract_variables(self, business_problem):
        """Extract relevant variables from business problem description"""
        if not self.client:
            # Fallback to mock data if no API key
            return self._get_mock_variables()
        
        try:
            system_prompt = """You are an expert business analyst. Extract relevant variables for marketing analysis from the given business problem. Return only a JSON array of variable names, nothing else."""
            
            user_prompt = f"""Business Problem: {business_problem}

Extract 6-10 relevant variables that would be important for analyzing this marketing/business problem. Focus on measurable, actionable variables.

Return format: ["Variable 1", "Variable 2", "Variable 3", ...]"""

            completion = self.client.chat.completions.create(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                model="llama-3.1-70b-versatile",
                temperature=0.7,
                max_tokens=1024
            )
            
            response = completion.choices[0].message.content.strip()
            variables = json.loads(response)
            return variables
            
        except Exception as e:
            print(f"Error extracting variables: {e}")
            return self._get_mock_variables()
    
    def _get_mock_variables(self):
        """Fallback mock variables"""
        return [
            "Customer Age",
            "Purchase Amount",
            "Product Category",
            "Marketing Channel",
            "Customer Location",
            "Purchase Frequency",
            "Customer Satisfaction Score",
            "Time to Purchase"
        ]