import json
from groq import Groq

class ChatInterface:
    def __init__(self):
        self.client = None
        self.conversation_history = []
    
    def set_api_key(self, api_key):
        """Set Groq API key"""
        self.client = Groq(api_key=api_key)
    
    def chat(self, message, project_data):
        """Chat with data using AI"""
        if not self.client:
            return self._get_mock_response(message)
        
        try:
            # Prepare data context
            data_context = self._prepare_data_context(project_data)
            
            system_prompt = """You are a data analyst assistant helping with marketing analysis. Answer questions about the user's data and provide insights. Be concise and actionable. Use the provided data context to give specific, data-driven responses."""
            
            user_prompt = f"""Data Context: {data_context}

User Question: {message}

Provide a helpful, data-driven response based on the available analysis results. Include specific numbers and insights where relevant."""

            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()
            
            # Store conversation
            self.conversation_history.append({
                'user': message,
                'assistant': response,
                'timestamp': pd.Timestamp.now()
            })
            
            return response
            
        except Exception as e:
            print(f"Error in chat: {e}")
            return self._get_mock_response(message)
    
    def _prepare_data_context(self, project_data):
        """Prepare data context for AI"""
        context = {}
        
        # Variables
        if 'variables' in project_data:
            context['variables'] = project_data['variables']
        
        # EDA Results
        if 'eda_results' in project_data:
            eda = project_data['eda_results']
            context['eda_summary'] = {
                'total_records': eda.get('summary', {}).get('total_records'),
                'correlations': len(eda.get('correlations', [])),
                'key_insights': eda.get('insights', [])[:3]  # Top 3 insights
            }
        
        # Model Results
        if 'model_results' in project_data:
            model = project_data['model_results']
            context['model_performance'] = {
                'accuracy': model.get('accuracy'),
                'model_type': model.get('model_type'),
                'top_features': model.get('feature_importance', [])[:3]
            }
        
        # Trend Results
        if 'trend_results' in project_data:
            trends = project_data['trend_results']
            context['trends'] = {
                'timeframe': trends.get('timeframe'),
                'key_trends': [t['metric'] + ': ' + t['direction'] for t in trends.get('trends', [])[:3]]
            }
        
        # Sentiment Results
        if 'sentiment_results' in project_data:
            sentiment = project_data['sentiment_results']
            context['sentiment'] = {
                'overall_positive': sentiment.get('overall', {}).get('positive'),
                'recommendations': sentiment.get('recommendations', [])[:2]
            }
        
        # A/B Test Results
        if 'ab_test_results' in project_data:
            ab_test = project_data['ab_test_results']
            context['ab_test'] = {
                'winner': ab_test.get('statistics', {}).get('winner'),
                'uplift': ab_test.get('statistics', {}).get('uplift'),
                'significance': ab_test.get('statistics', {}).get('significance')
            }
        
        return json.dumps(context, indent=2)
    
    def _get_mock_response(self, message):
        """Generate mock response when AI is not available"""
        message_lower = message.lower()
        
        if 'customer' in message_lower and 'satisfaction' in message_lower:
            return "Based on your sentiment analysis, customer satisfaction is at 68% positive. Key drivers include product quality and delivery speed. Consider focusing on the 10% negative feedback to improve overall satisfaction."
        
        elif 'marketing' in message_lower and 'channel' in message_lower:
            return "Your data shows that Email and Social Media are the top-performing marketing channels with higher conversion rates. TV and Print show lower engagement. Consider reallocating budget to digital channels for better ROI."
        
        elif 'revenue' in message_lower or 'forecast' in message_lower:
            return "Based on trend analysis, revenue is projected to grow by 15.3% next quarter. The forecast shows $267.3K with 78% confidence. Key growth drivers include customer acquisition and increased purchase frequency."
        
        elif 'correlation' in message_lower or 'relationship' in message_lower:
            return "Strong correlations detected between Customer Age and Purchase Amount (0.65), and between Satisfaction Score and Purchase Frequency (0.58). These relationships suggest targeting strategies based on age demographics."
        
        elif 'segment' in message_lower or 'group' in message_lower:
            return "Your predictive model identifies three customer segments: High-value (73%), Medium-value (21%), and Low-value (6%). Focus retention efforts on high-value customers and conversion strategies for medium-value segments."
        
        elif 'test' in message_lower or 'experiment' in message_lower:
            return "Your A/B test shows the treatment variant performs 21.1% better than control with statistical significance (p=0.023). Recommend implementing the treatment for your full campaign to capture the revenue uplift."
        
        else:
            return "I can help you analyze your marketing data. Ask me about customer segments, marketing channel performance, revenue forecasts, correlations, or A/B test results. What specific insights would you like to explore?"
    
    def get_conversation_history(self):
        """Get chat conversation history"""
        return self.conversation_history
    
    def clear_history(self):
        """Clear conversation history"""
        self.conversation_history = []
        return "Conversation history cleared."