"""
Integrated Workforce Optimization System
Combines intent parsing with skill-based optimization
"""

import json
from intent_parser import IntentParser
from skill_optimizer import SkillOptimizer
from typing import Dict, Any, Optional
from openai import OpenAI

class IntegratedOptimizer:
    def __init__(self, api_key=None):
        """Initialize both components"""
        self.intent_parser = IntentParser(api_key=api_key)
        self.skill_optimizer = SkillOptimizer()
        self.client = OpenAI(api_key=api_key) if api_key else None
        
    def process_optimization_request(self, user_input: str, algorithm='balanced') -> Dict[str, Any]:
        """
        Complete pipeline from natural language to optimization results
        
        Args:
            user_input: Natural language request
            algorithm: Optimization algorithm to use
            
        Returns:
            Complete optimization results
        """
        print("\n" + "="*70)
        print("INTEGRATED WORKFORCE OPTIMIZATION")
        print("="*70)
        
        # Step 1: Parse intent and adjust skills
        print("\n📝 STEP 1: Understanding Request")
        print("-"*40)
        parsed = self.intent_parser.process_user_request(user_input)
        
        # Extract reduction target
        reduction_target = parsed.get('reduction_target')
        if not reduction_target:
            print("❌ No reduction target found in request")
            return None
        
        reduction_percentage = reduction_target.get('value', 0)
        
        # Use adjusted skills as requirements
        required_skills = parsed.get('adjusted_skills', 
                                    self.intent_parser.default_skills)
        
        # Step 2: Run optimization
        print("\n🔧 STEP 2: Running Optimization")
        print("-"*40)
        
        optimization_result = self.skill_optimizer.optimize(
            reduction_percentage=reduction_percentage,
            required_skills=required_skills,
            algorithm=algorithm
        )
        
        # Step 3: Combine results
        combined_result = {
            'request': {
                'original_input': user_input,
                'reduction_target': reduction_target,
                'intent': parsed.get('intent'),
                'skill_adjustments': parsed.get('skill_changes')
            },
            'optimization': optimization_result,
            'adjusted_skills': required_skills,
            'recommendations': self._generate_recommendations(
                optimization_result, 
                parsed.get('intent', {})
            )
        }
        
        # Step 4: Generate natural language explanation
        if self.client:
            combined_result['explanation'] = self.generate_natural_explanation(combined_result)
        
        return combined_result
    
    def _generate_recommendations(self, optimization_result: Dict, intent: Dict) -> list:
        """Generate actionable recommendations based on results"""
        recommendations = []
        
        # Check for critical skill gaps
        if optimization_result['analysis']['critical_impacts']:
            recommendations.append({
                'type': 'warning',
                'message': f"Critical skill gaps identified: {len(optimization_result['analysis']['critical_impacts'])} skills below required levels",
                'action': 'Consider retraining or selective retention'
            })
        
        # Cost savings
        saved = optimization_result['analysis']['cost']['saved']
        recommendations.append({
            'type': 'info',
            'message': f'Annual cost savings: ${saved:,.0f}',
            'action': 'Allocate portion to retraining programs'
        })
        
        # Intent-specific recommendations
        if intent.get('detected'):
            focus_areas = intent.get('focus_areas', [])
            if 'artificial intelligence' in focus_areas or 'machine learning' in focus_areas:
                recommendations.append({
                    'type': 'suggestion',
                    'message': 'AI/ML focus detected',
                    'action': 'Prioritize retention of ML engineers and data scientists'
                })
            elif 'backend' in ' '.join(focus_areas).lower():
                recommendations.append({
                    'type': 'suggestion',
                    'message': 'Backend focus detected',
                    'action': 'Consider cross-training frontend engineers for backend roles'
                })
        
        return recommendations
    
    def generate_natural_explanation(self, result: Dict[str, Any]) -> str:
        """Generate a natural language explanation of the optimization results"""
        if not self.client:
            return "Optimization complete. See details above."
        
        # Prepare context for the LLM
        opt = result['optimization']
        req = result['request']
        
        # Prepare removed employees summary
        removed_employees = []
        for emp in opt['removed_employees'][:5]:  # Top 5 for context
            removed_employees.append(f"- {emp['name']} ({emp['role']}): ${emp['salary']:,}")
        
        # Prepare skill changes summary  
        skill_changes = []
        for skill, data in sorted(opt['analysis']['skill_changes'].items(), 
                                key=lambda x: abs(x[1]['change']), 
                                reverse=True)[:5]:
            skill_name = skill.replace('_', ' ').title()
            change = data['change']
            symbol = '↑' if change > 0 else '↓'
            skill_changes.append(f"- {skill_name}: {data['before']:.1f} → {data['after']:.1f} ({symbol}{abs(change):.1f})")
        
        prompt = f"""You are an HR optimization assistant explaining the results of a workforce optimization.
        
User Request: {req['original_input']}

Optimization Context:
- Team Size: {opt['analysis']['team_size']['before']} → {opt['analysis']['team_size']['after']} employees
- Employees Removed: {opt['analysis']['team_size']['reduced']}
- Cost Saved: ${opt['analysis']['cost']['saved']:,.0f} annually
- Strategic Intent: {req['intent'].get('description', 'General optimization')}
- Focus Areas: {', '.join(req['intent'].get('focus_areas', ['balanced optimization']))}
- Algorithm Used: {opt.get('algorithm', 'balanced')}

Key Employees Removed (sample):
{chr(10).join(removed_employees)}

Top Skill Level Changes:
{chr(10).join(skill_changes)}

Critical Skill Gaps: {len(opt['analysis'].get('critical_impacts', []))}
Skill Gap Score: {opt['final_gap']:.2f}

Generate a natural, conversational explanation that:
1. Acknowledges the user's request and strategic intent
2. Explains WHY certain people were removed (based on skills, redundancy, strategic alignment)
3. Highlights the key outcomes (team size, cost savings, skill impacts)
4. Mentions any critical concerns or risks
5. Provides forward-looking advice
6. Uses a professional but friendly tone
7. Be specific about the optimization logic and decisions made

Keep the response under 300 words and format it with clear sections using markdown."""

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "system", "content": "You are an expert HR optimization assistant providing clear, actionable insights."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=500
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            print(f"Error generating explanation: {e}")
            return self._generate_fallback_explanation(result)
    
    def _generate_fallback_explanation(self, result: Dict[str, Any]) -> str:
        """Generate a fallback explanation if API fails"""
        opt = result['optimization']
        req = result['request']
        
        explanation = f"""## Optimization Complete ✅

I've successfully optimized your engineering team based on your request to {req['original_input'].lower()}.

**Key Results:**
- Team reduced from {opt['analysis']['team_size']['before']} to {opt['analysis']['team_size']['after']} employees
- Annual cost savings: ${opt['analysis']['cost']['saved']:,.0f}
- {opt['analysis']['team_size']['reduced']} employees removed

**Strategic Alignment:**
The optimization focused on {req['intent'].get('description', 'maintaining balanced skill coverage').lower()}.

**Skill Impact:**
The team's average skill levels have been adjusted to align with your strategic priorities. 
"""
        
        if opt['analysis'].get('critical_impacts'):
            explanation += f"\n⚠️ **Note:** {len(opt['analysis']['critical_impacts'])} skills are below required levels. Consider targeted training or selective retention.\n"
        else:
            explanation += "\n✅ All critical skills maintained above required levels.\n"
        
        return explanation
    
    def print_integrated_results(self, result: Dict[str, Any]):
        """Pretty print integrated results"""
        if not result:
            return
        
        print("\n" + "="*70)
        print("OPTIMIZATION COMPLETE")
        print("="*70)
        
        # Request summary
        req = result['request']
        print(f"\n📋 REQUEST SUMMARY:")
        print(f"  Input: {req['original_input'][:80]}...")
        if req['intent'].get('detected'):
            print(f"  Intent: {req['intent']['description']}")
            print(f"  Focus: {', '.join(req['intent'].get('focus_areas', []))}")
        
        # Optimization results
        self.skill_optimizer.print_optimization_results(result['optimization'])
        
        # Recommendations
        print(f"\n{'='*60}")
        print("💡 RECOMMENDATIONS:")
        print(f"{'='*60}")
        for rec in result['recommendations']:
            emoji = {'warning': '⚠️', 'info': 'ℹ️', 'suggestion': '💡'}.get(rec['type'], '•')
            print(f"\n{emoji} {rec['message']}")
            print(f"   → {rec['action']}")
    
    def compare_algorithms(self, user_input: str) -> Dict[str, Any]:
        """Compare greedy vs balanced algorithms"""
        print("\n" + "="*70)
        print("ALGORITHM COMPARISON")
        print("="*70)
        
        # Parse once
        parsed = self.intent_parser.process_user_request(user_input)
        reduction_target = parsed.get('reduction_target')
        
        if not reduction_target:
            print("❌ No reduction target found")
            return None
        
        reduction_percentage = reduction_target.get('value', 0)
        required_skills = parsed.get('adjusted_skills', self.intent_parser.default_skills)
        
        # Run both algorithms
        results = {}
        for algo in ['greedy', 'balanced']:
            print(f"\n🔄 Running {algo} algorithm...")
            results[algo] = self.skill_optimizer.optimize(
                reduction_percentage=reduction_percentage,
                required_skills=required_skills,
                algorithm=algo
            )
        
        # Compare results
        print("\n" + "="*70)
        print("COMPARISON RESULTS")
        print("="*70)
        
        print(f"\n{'Metric':<30} {'Greedy':>15} {'Balanced':>15}")
        print("-"*60)
        
        for algo in ['greedy', 'balanced']:
            r = results[algo]
            if algo == 'greedy':
                print(f"{'Skill Gap Score':<30} {r['final_gap']:>15.2f}", end='')
            else:
                print(f" {r['final_gap']:>15.2f}")
        
        for algo in ['greedy', 'balanced']:
            r = results[algo]
            if algo == 'greedy':
                print(f"{'Cost Saved':<30} ${r['analysis']['cost']['saved']:>14,.0f}", end='')
            else:
                print(f" ${r['analysis']['cost']['saved']:>14,.0f}")
        
        for algo in ['greedy', 'balanced']:
            r = results[algo]
            critical = len(r['analysis']['critical_impacts'])
            if algo == 'greedy':
                print(f"{'Critical Skill Gaps':<30} {critical:>15}", end='')
            else:
                print(f" {critical:>15}")
        
        # Recommendation
        if results['greedy']['final_gap'] < results['balanced']['final_gap']:
            print("\n\n✅ Recommendation: Use GREEDY algorithm (lower skill gap)")
        else:
            print("\n\n✅ Recommendation: Use BALANCED algorithm (better overall)")
        
        return results


def main():
    """Test integrated system"""
    
    # Test cases
    test_inputs = [
        "Reduce engineering by 20% and focus heavily on AI and machine learning capabilities",
        "Cut 25% of the team while shutting down frontend to focus on backend APIs",
        "Optimize team by 15% to become a data platform company"
    ]
    
    # Use your API key
    api_key = "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"
    
    optimizer = IntegratedOptimizer(api_key=api_key)
    
    # Test first case
    print("\n" + "🚀 "*20)
    print("TESTING INTEGRATED OPTIMIZATION")
    print("🚀 "*20)
    
    for test_input in test_inputs[:1]:  # Test just the first one
        result = optimizer.process_optimization_request(test_input)
        optimizer.print_integrated_results(result)
        
        # Save detailed results
        with open('integrated_result.json', 'w') as f:
            # Make it JSON serializable
            output = {
                'request': result['request'],
                'team_changes': {
                    'removed': [e['name'] for e in result['optimization']['removed_employees']],
                    'remaining_count': len(result['optimization']['remaining_employees']),
                    'cost_saved': result['optimization']['analysis']['cost']['saved']
                },
                'recommendations': result['recommendations']
            }
            json.dump(output, f, indent=2)
        
        print("\n💾 Detailed results saved to integrated_result.json")
    
    # Compare algorithms
    print("\n\n" + "="*70)
    print("ALGORITHM COMPARISON TEST")
    print("="*70)
    optimizer.compare_algorithms(test_inputs[0])


if __name__ == "__main__":
    main()