"""
Intent Parser for Workforce Optimization
Extracts headcount reduction targets and adjusts skill requirements based on strategic intent
"""

import json
from typing import Dict, Any, Optional, Tuple
from openai import OpenAI
import os
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from typing import List
from typing import Literal, Union

# Load environment variables
load_dotenv()

# Default engineering skills and required levels
DEFAULT_ENGINEERING_SKILLS = {
    # Technical Skills
    "python": 3.0,
    "javascript": 2.5,
    "react": 2.5,
    "backend_design": 3.0,
    "system_architecture": 3.0,
    "api_development": 3.0,
    "database_management": 2.5,
    "devops": 2.5,
    "cloud_infrastructure": 2.5,
    "security": 2.5,
    
    # AI/ML Skills
    "machine_learning": 2.0,
    "pytorch": 2.0,
    "tensorflow": 1.5,
    "data_engineering": 2.5,
    "data_analysis": 2.5,
    
    # Soft Skills
    "communication": 3.0,
    "problem_solving": 3.5,
    "team_collaboration": 3.0,
    "leadership": 2.0,
    "project_management": 2.0,
    
    # Specialized Skills
    "frontend_optimization": 2.0,
    "mobile_development": 1.5,
    "testing": 2.5,
    "documentation": 2.5,
    "code_review": 3.0
}

# Prompt for extracting reduction target and intent
EXTRACTION_PROMPT = """You are an expert at understanding workforce optimization requests.

Given a user's input, extract TWO things:
1. The headcount reduction target (either as a percentage or absolute number)
2. The strategic intent or initiative (what the company is focusing on)

Return a JSON with this exact structure:
{
    "reduction_target": {
        "type": "percentage" or "absolute",
        "value": <number>
    },
    "intent": {
        "detected": true/false,
        "description": "<brief description of the strategic intent>",
        "focus_areas": ["<area1>", "<area2>", ...],
        "de_emphasis_areas": ["<area1>", "<area2>", ...]
    }
}

Examples:
- Input: "Reduce headcount by 20%"
  Output: {"reduction_target": {"type": "percentage", "value": 20}, "intent": {"detected": false}}

- Input: "Cut 5 engineers and focus on AI initiatives"
  Output: {"reduction_target": {"type": "absolute", "value": 5}, "intent": {"detected": true, "description": "Focus on AI initiatives", "focus_areas": ["artificial intelligence", "machine learning"], "de_emphasis_areas": []}}

- Input: "We need to reduce costs by 15% and shift from web to mobile development"
  Output: {"reduction_target": {"type": "percentage", "value": 15}, "intent": {"detected": true, "description": "Shift from web to mobile development", "focus_areas": ["mobile development"], "de_emphasis_areas": ["web development"]}}

User Input: {user_input}

Return ONLY valid JSON, no additional text."""

# Prompt for adjusting skill requirements based on intent
SKILL_ADJUSTMENT_PROMPT = """You are an expert at translating strategic business initiatives into technical skill requirements.

Given:
1. Current skill requirements for an engineering team
2. A strategic intent/initiative

Adjust the skill requirement levels to align with the strategic direction.

Rules:
- You can ONLY adjust the levels of existing skills (range: 0.0 to 5.0)
- You CANNOT add new skills
- Increase levels for skills critical to the initiative (max increase: +2.0)
- Decrease levels for skills less relevant to the initiative (max decrease: -2.0)
- Most adjustments should be in the 0.5-1.0 range
- Keep at least a minimum level (1.0) for basic skills even if de-emphasized

Current skill requirements:
{current_skills}

Strategic Intent: {intent_description}
Focus Areas: {focus_areas}
De-emphasis Areas: {de_emphasis_areas}

Return ONLY a JSON object with the SAME skill names and adjusted levels.
Provide brief reasoning for significant changes (>1.0 change).

Example output format:
{
    "adjusted_skills": {
        "python": 3.5,
        "javascript": 2.0,
        ...all skills must be included...
    },
    "adjustments_reasoning": {
        "python": "Increased by 0.5 as it's critical for the initiative",
        "javascript": "Decreased by 0.5 as frontend is de-emphasized"
    }
}"""

class ReductionTarget(BaseModel):
    type: Union[Literal["percentage"], Literal["absolute"]] = Field(description="The type of reduction target")
    value: float = Field(description="The value of the reduction target")

class Intent(BaseModel):
    detected: bool = Field(description="Whether the intent was detected")
    description: str = Field(description="The description of the intent")
    focus_areas: List[str] = Field(description="The areas of focus for the intent")
    de_emphasis_areas: List[str] = Field(description="The areas of de-emphasis for the intent")


class ExtactionUnderstanding(BaseModel):
    reduction_target: ReductionTarget = Field(description="The reduction target")
    intent: Intent = Field(description="The intent")

class AdjustedSkills(BaseModel):
    skill: str = Field(description="The skill name")
    level: float = Field(description="The level of the skill")

class AdjustmentsReasoning(BaseModel):
    skill: str = Field(description="The skill name")
    reasoning: str = Field(description="The reasoning for the adjustment")


class SkillAdjustment(BaseModel):
    adjusted_skills: List[AdjustedSkills] = Field(description="The adjusted skill levels")
    adjustments_reasoning: List[AdjustmentsReasoning] = Field(description="The reasoning for the adjustments")


class IntentParser:
    def __init__(self, api_key: Optional[str] = None):
        """Initialize the Intent Parser with OpenAI client"""
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError("OpenAI API key not found. Set OPENAI_API_KEY environment variable.")
        
        self.client = OpenAI(api_key=self.api_key)
        self.default_skills = DEFAULT_ENGINEERING_SKILLS.copy()
    
    def extract_intent_and_target(self, user_input: str) -> Dict[str, Any]:
        """
        Extract headcount reduction target and strategic intent from user input
        
        Args:
            user_input: Natural language input from user
            
        Returns:
            Dictionary with reduction target and intent information
        """
        prompt = EXTRACTION_PROMPT.replace("{user_input}", user_input)
        
        try:
            response = self.client.responses.parse(
                model="gpt-4o",
                input=[
                    {"role": "system", "content": "You are a precise JSON generator. Only output valid JSON."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.1,  # Low temperature for consistent extraction
                text_format=ExtactionUnderstanding
            )

            parsed: ExtactionUnderstanding = response.output_parsed
            return parsed.model_dump()
            
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON response: {e}")
            return {
                "reduction_target": None,
                "intent": {"detected": False},
                "error": "Failed to parse response"
            }
        except Exception as e:
            print(f"Error calling OpenAI API: {e}")
            return {
                "reduction_target": None,
                "intent": {"detected": False},
                "error": str(e)
            }
    
    def adjust_skills_for_intent(self, intent_info: Dict[str, Any]) -> Dict[str, float]:
        """
        Adjust skill requirements based on detected intent
        
        Args:
            intent_info: Intent information from extraction
            
        Returns:
            Dictionary with adjusted skill levels
        """
        # If no intent detected, return default skills
        if not intent_info.get("detected", False):
            return self.default_skills.copy()
        
        # Prepare the prompt
        # prompt = SKILL_ADJUSTMENT_PROMPT.format(
        #     current_skills=json.dumps(self.default_skills, indent=2),
        #     intent_description=intent_info.get("description", ""),
        #     focus_areas=json.dumps(intent_info.get("focus_areas", [])),
        #     de_emphasis_areas=json.dumps(intent_info.get("de_emphasis_areas", []))

        prompt = SKILL_ADJUSTMENT_PROMPT.replace("{current_skills}", json.dumps(self.default_skills, indent=2))
        prompt = prompt.replace("{intent_description}", intent_info.get("description", ""))
        prompt = prompt.replace("{focus_areas}", json.dumps(intent_info.get("focus_areas", [])))
        prompt = prompt.replace("{de_emphasis_areas}", json.dumps(intent_info.get("de_emphasis_areas", [])))
        
        try:
            response = self.client.responses.parse(
                model="gpt-4o",
                input=[
                    {"role": "system", "content": "You are an expert at workforce skill optimization. Adjust skill levels based on strategic initiatives."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.2,
                text_format=SkillAdjustment
            )
            
            result: SkillAdjustment = response.output_parsed
            result = result.model_dump()
            adjusted_skills = result.get("adjusted_skills", self.default_skills)
            adjusted_skills = {skill['skill']: skill['level'] for skill in adjusted_skills}
            print(f"Adjusted skills: {adjusted_skills}")
            
            # Validate that all skills are present and within range
            validated_skills = {}
            for skill in self.default_skills.keys():
                if skill in adjusted_skills:
                    # Ensure skill level is within valid range
                    level = max(0.0, min(5.0, float(adjusted_skills[skill])))
                    validated_skills[skill] = level
                else:
                    validated_skills[skill] = self.default_skills[skill]
            result["adjustments_reasoning"] = {skill['skill']: skill['reasoning'] for skill in result["adjustments_reasoning"]}
            # Print reasoning if available
            if "adjustments_reasoning" in result:
                print("\nAdjustment Reasoning:")
                for skill, reason in result["adjustments_reasoning"].items():
                    print(f"  - {skill}: {reason}")
            
            return validated_skills
            
        except Exception as e:
            print(f"Error adjusting skills: {e}")
            return self.default_skills.copy()
    
    def process_user_request(self, user_input: str) -> Dict[str, Any]:
        """
        Complete pipeline: extract intent and adjust skills
        
        Args:
            user_input: Natural language input from user
            
        Returns:
            Complete optimization parameters
        """
        print(f"\n{'='*60}")
        print(f"Processing: {user_input}")
        print(f"{'='*60}")
        
        # Step 1: Extract intent and target
        print("\n1. Extracting intent and target...")
        extraction_result = self.extract_intent_and_target(user_input)
        
        print(f"\nExtracted:")
        print(f"  Reduction Target: {extraction_result.get('reduction_target')}")
        print(f"  Intent Detected: {extraction_result.get('intent', {}).get('detected')}")
        if extraction_result.get('intent', {}).get('detected'):
            print(f"  Intent: {extraction_result.get('intent', {}).get('description')}")
            print(f"  Focus Areas: {extraction_result.get('intent', {}).get('focus_areas')}")
            print(f"  De-emphasis Areas: {extraction_result.get('intent', {}).get('de_emphasis_areas')}")
        
        # Step 2: Adjust skills based on intent
        print("\n2. Adjusting skill requirements...")
        adjusted_skills = self.adjust_skills_for_intent(extraction_result.get('intent', {}))
        
        # Show significant changes
        print("\nSignificant Skill Changes:")
        for skill, new_level in adjusted_skills.items():
            old_level = self.default_skills[skill]
            change = new_level - old_level
            if abs(change) >= 0.5:
                direction = "↑" if change > 0 else "↓"
                print(f"  {skill}: {old_level:.1f} → {new_level:.1f} ({direction}{abs(change):.1f})")
        
        return {
            "user_input": user_input,
            "reduction_target": extraction_result.get('reduction_target'),
            "intent": extraction_result.get('intent'),
            "default_skills": self.default_skills,
            "adjusted_skills": adjusted_skills,
            "skill_changes": {
                skill: adjusted_skills[skill] - self.default_skills[skill]
                for skill in self.default_skills.keys()
            }
        }


def main():
    """Test the intent parser with various examples"""
    
    # Initialize parser
    parser = IntentParser()
    
    # Test cases
    test_cases = [
        "Reduce engineering headcount by 20%",
        "Cut 15% of the team and focus on AI and machine learning capabilities",
        "We need to reduce costs by 25%. We're shutting down mobile development and focusing entirely on backend infrastructure",
        "Reduce headcount by 30% while transitioning from a product company to a services company",
        "Cut 10 engineers. We're moving all development offshore except for architecture and security",
        "Optimize team size by 15% and double down on our data platform initiatives"
    ]
    
    results = []
    for test_input in test_cases:
        result = parser.process_user_request(test_input)
        results.append(result)
        print("\n" + "="*60)
    
    # Save results to file for review
    with open("intent_parsing_results.json", "w") as f:
        json.dump(results, f, indent=2)
    
    print("\n\nResults saved to intent_parsing_results.json")


if __name__ == "__main__":
    # For testing without API key, use mock mode
    import sys
    if "--mock" in sys.argv:
        print("Running in mock mode (no actual API calls)")
        # Add mock implementation here if needed
    else:
        main()