import sys
import json
import requests
from typing import List, Dict, Any, Optional
from dataclasses import dataclass, asdict
import numpy as np
from collections import defaultdict

@dataclass
class CourseRecommendation:
    courseId: str
    courseName: str
    matchScore: str  # "85%" format
    reason: str
    relevance_factors: Dict[str, Any]

class AdvancedCourseRecommender:
    """
    Professional course recommendation system using Llama 3.1 8B via Ollama.
    Provides intelligent, context-aware recommendations with detailed reasoning.
    """
    
    def __init__(self, ollama_url: str = "http://localhost:11434"):
        self.ollama_url = ollama_url
        self.model = "llama3.1:8b"  # Fast and excellent for analysis
        self.token_count = 0
        
    def generate_recommendations(
        self, 
        user_data: Dict[str, Any],
        users_under_manager: Dict[str, Any],
        all_courses: List[Dict[str, Any]],
        logged_in_user_id: str
    ) -> Dict[str, Any]:
        """
        Generate top 5 course recommendations with advanced analysis.
        """
        
        # Extract logged-in user info
        logged_user = user_data.get(logged_in_user_id, {})
        
        # Analyze team context if manager
        team_context = self._analyze_team_learning(
            logged_user, 
            users_under_manager
        )
        
        # Get user's current state
        user_skills = self._extract_skills(logged_user.get("skills", []))
        completed_courses = set(logged_user.get("completedCourses", []))
        assigned_courses = set(logged_user.get("assignedCourses", []))
        
        # Filter available courses
        available_courses = [
            course for course in all_courses 
            if course["courseId"] not in completed_courses
        ]
        
        if not available_courses:
            return self._create_response([], "No new courses available.")
        
        # Analyze each course and generate recommendations
        recommendations = []
        
        for course in available_courses:
            analysis = self._analyze_course_fit(
                course,
                logged_user,
                team_context,
                user_skills,
                assigned_courses
            )
            
            if analysis:
                recommendations.append(analysis)
        
        # Sort by score and get top 5
        recommendations.sort(key=lambda x: x["raw_score"], reverse=True)
        top_recommendations = recommendations[:5]
        
        # Generate detailed reasons using LLM
        final_recommendations = []
        for rec in top_recommendations:
            detailed_rec = self._generate_detailed_recommendation(
                rec,
                logged_user,
                team_context
            )
            final_recommendations.append(detailed_rec)
        
        return self._create_response(
            final_recommendations,
            f"Generated {len(final_recommendations)} personalized course recommendations."
        )
    
    def _analyze_team_learning(
        self, 
        logged_user: Dict[str, Any],
        users_under_manager: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Analyze what courses team members have completed that the user hasn't.
        """
        team_context = {
            "popular_completed_courses": [],
            "team_skills": set(),
            "team_size": 0,
            "missing_important_courses": []
        }
        
        users = users_under_manager.get("users", {})
        if not users:
            return team_context
        
        team_context["team_size"] = len(users)
        user_completed = set(logged_user.get("completedCourses", []))
        
        # Track course completion frequency
        course_completion_count = defaultdict(int)
        
        for user_id, user_info in users.items():
            if user_id != str(logged_user.get("user_id", "")):
                completed = user_info.get("completedCourses", [])
                for course_id in completed:
                    course_completion_count[course_id] += 1
                    
                # Collect team skills
                for skill in user_info.get("skills", []):
                    team_context["team_skills"].add(skill.get("skill_name", ""))
        
        # Find popular courses not completed by logged-in user
        for course_id, count in course_completion_count.items():
            if course_id not in user_completed:
                completion_rate = (count / team_context["team_size"]) * 100
                if completion_rate >= 50:  # 50% or more of team completed
                    team_context["missing_important_courses"].append({
                        "course_id": course_id,
                        "completion_rate": completion_rate
                    })
        
        return team_context
    
    def _extract_skills(self, skills_list: List[Dict[str, Any]]) -> Dict[str, str]:
        """Extract skills categorized by type."""
        skills = {
            "technical": [],
            "domain": [],
            "soft": []
        }
        
        for skill in skills_list:
            skill_name = skill.get("skill_name", "")
            skill_type = skill.get("skill_type", "")
            
            if skill_type == "1":
                skills["technical"].append(skill_name)
            elif skill_type == "2":
                skills["domain"].append(skill_name)
            elif skill_type == "3":
                skills["soft"].append(skill_name)
        
        return skills
    
    def _analyze_course_fit(
        self,
        course: Dict[str, Any],
        user: Dict[str, Any],
        team_context: Dict[str, Any],
        user_skills: Dict[str, str],
        assigned_courses: set
    ) -> Optional[Dict[str, Any]]:
        """
        Analyze how well a course fits the user's needs.
        Returns analysis with scoring factors.
        """
        course_id = course["courseId"]
        course_name = course["name"]
        course_skills = [s.get("skill_name", "") for s in course.get("skills", [])]
        
        # Scoring factors (0-100 each)
        scores = {
            "skill_alignment": 0,
            "team_relevance": 0,
            "career_progression": 0,
            "priority": 0
        }
        
        # 1. Skill Alignment (40% weight)
        if course_skills:
            user_skill_names = (
                user_skills["technical"] + 
                user_skills["domain"] + 
                user_skills["soft"]
            )
            
            # Check for skill gaps (courses teaching new skills)
            new_skills = [s for s in course_skills if s not in user_skill_names]
            related_skills = [s for s in course_skills if s in user_skill_names]
            
            if new_skills:
                scores["skill_alignment"] = 80  # New skills to learn
            elif related_skills:
                scores["skill_alignment"] = 60  # Advancing existing skills
            else:
                scores["skill_alignment"] = 40  # General learning
        else:
            scores["skill_alignment"] = 50  # No specific skills listed
        
        # 2. Team Relevance (30% weight)
        if team_context["missing_important_courses"]:
            for missing_course in team_context["missing_important_courses"]:
                if missing_course["course_id"] == course_id:
                    scores["team_relevance"] = missing_course["completion_rate"]
                    break
        
        # 3. Career Progression (20% weight)
        user_designation = user.get("designation", "").lower()
        course_description = (
            course.get("description", "") + 
            course.get("short_description", "")
        ).lower()
        
        # Keywords indicating advanced/leadership content
        advanced_keywords = [
            "advanced", "leadership", "management", "strategy", 
            "architect", "senior", "expert", "professional"
        ]
        
        if any(keyword in course_description for keyword in advanced_keywords):
            scores["career_progression"] = 75
        else:
            scores["career_progression"] = 50
        
        # 4. Priority (10% weight)
        if course_id in assigned_courses:
            scores["priority"] = 100  # Assigned by manager
        else:
            scores["priority"] = 50
        
        # Calculate weighted score
        raw_score = (
            scores["skill_alignment"] * 0.40 +
            scores["team_relevance"] * 0.30 +
            scores["career_progression"] * 0.20 +
            scores["priority"] * 0.10
        )
        
        return {
            "courseId": course_id,
            "courseName": course_name,
            "raw_score": raw_score,
            "scores": scores,
            "course_skills": course_skills,
            "course_description": course.get("short_description", "")
        }
    
    def _generate_detailed_recommendation(
        self,
        course_analysis: Dict[str, Any],
        user: Dict[str, Any],
        team_context: Dict[str, Any]
    ) -> CourseRecommendation:
        """
        Use LLM to generate detailed reasoning for recommendation.
        """
        # Prepare context for LLM
        user_name = user.get("user_name", "User")
        user_skills_text = self._format_skills_for_llm(user.get("skills", []))
        course_name = course_analysis["courseName"]
        course_skills = ", ".join(course_analysis["course_skills"]) if course_analysis["course_skills"] else "General skills"
        match_score = int(course_analysis["raw_score"])
        
        # Build context
        context_parts = []
        
        if course_analysis["scores"]["team_relevance"] > 0:
            context_parts.append(
                f"{int(course_analysis['scores']['team_relevance'])}% of your team has completed this course"
            )
        
        if course_analysis["scores"]["priority"] == 100:
            context_parts.append("This course has been assigned to you by your manager")
        
        context = ". ".join(context_parts) if context_parts else "Based on your profile"
        
        # Create prompt for LLM
        prompt = f"""You are a professional learning advisor. Generate a concise, compelling reason (2-3 sentences) why this course is recommended.

User: {user_name}
Current Skills: {user_skills_text}
Recommended Course: {course_name}
Course Teaches: {course_skills}
Match Score: {match_score}%
Context: {context}

Provide a professional, specific reason focusing on:
1. How this course fills skill gaps or advances current skills
2. Career impact and practical applications
3. Relevance to team/organizational needs (if applicable)

Reason (2-3 sentences):"""

        try:
            reason = self._call_ollama(prompt, max_tokens=150)
        except Exception as e:
            # Fallback reasoning
            reason = self._generate_fallback_reason(course_analysis, user, context)
        
        return CourseRecommendation(
            courseId=course_analysis["courseId"],
            courseName=course_analysis["courseName"],
            matchScore=f"{match_score}%",
            reason=reason.strip(),
            relevance_factors={
                "skill_alignment": f"{int(course_analysis['scores']['skill_alignment'])}%",
                "team_relevance": f"{int(course_analysis['scores']['team_relevance'])}%",
                "career_progression": f"{int(course_analysis['scores']['career_progression'])}%",
                "priority": f"{int(course_analysis['scores']['priority'])}%"
            }
        )
    
    def _format_skills_for_llm(self, skills: List[Dict[str, Any]]) -> str:
        """Format skills for LLM context."""
        if not skills:
            return "No specific skills listed"
        
        skill_names = [s.get("skill_name", "") for s in skills if s.get("skill_name")]
        return ", ".join(skill_names) if skill_names else "General skills"
    
    def _generate_fallback_reason(
        self,
        course_analysis: Dict[str, Any],
        user: Dict[str, Any],
        context: str
    ) -> str:
        """Generate reasoning without LLM if API fails."""
        course_name = course_analysis["courseName"]
        match_score = int(course_analysis["raw_score"])
        
        reason_parts = []
        
        if course_analysis["scores"]["team_relevance"] > 60:
            reason_parts.append(
                f"This course is highly relevant as majority of your team has completed it"
            )
        
        if course_analysis["course_skills"]:
            skills_text = ", ".join(course_analysis["course_skills"][:3])
            reason_parts.append(
                f"'{course_name}' will help you develop skills in {skills_text}"
            )
        else:
            reason_parts.append(
                f"'{course_name}' provides valuable learning opportunities aligned with your career growth"
            )
        
        if course_analysis["scores"]["priority"] == 100:
            reason_parts.append("This is a priority course assigned by your manager")
        
        return ". ".join(reason_parts) + "."
    
    def _call_ollama(self, prompt: str, max_tokens: int = 150) -> str:
        """
        Call Ollama API for text generation.
        """
        try:
            response = requests.post(
                f"{self.ollama_url}/api/generate",
                json={
                    "model": self.model,
                    "prompt": prompt,
                    "stream": False,
                    "options": {
                        "temperature": 0.7,
                        "top_p": 0.9,
                        "num_predict": max_tokens
                    }
                },
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                self.token_count += result.get("eval_count", 0)
                return result.get("response", "").strip()
            else:
                raise Exception(f"Ollama API error: {response.status_code}")
                
        except Exception as e:
            print(f"Ollama API call failed: {str(e)}", file=sys.stderr)
            raise
    
    def _create_response(
        self,
        recommendations: List[CourseRecommendation],
        message: str
    ) -> Dict[str, Any]:
        """Create final response structure."""
        return {
            "recommendations": [asdict(rec) for rec in recommendations],
            "total_count": len(recommendations),
            "message": message,
            "token_count": self.token_count
        }


def main():
    """Main entry point for the recommendation system."""
    if len(sys.argv) < 2:
        print(json.dumps({
            "error": "No input file provided",
            "recommendations": [],
            "token_count": 0
        }))
        sys.exit(1)
    
    try:
        # Load input data
        with open(sys.argv[1], 'r') as f:
            data = json.load(f)
        
        client_id = str(data.get("client_id", ""))
        user_data = data.get("user_data", {})
        users_under_manager = data.get("users_under_manager_data", {})
        all_courses_data = data.get("client_all_courses_data", {})
        
        # Get courses for this client
        all_courses = all_courses_data.get(client_id, [])
        
        if not all_courses:
            print(json.dumps({
                "recommendations": [],
                "total_count": 0,
                "message": "No courses available for this client",
                "token_count": 0
            }))
            sys.exit(0)
        
        # Get logged-in user ID (first user in user_data)
        logged_in_user_id = list(user_data.keys())[0] if user_data else None
        
        if not logged_in_user_id:
            print(json.dumps({
                "error": "No user data provided",
                "recommendations": [],
                "token_count": 0
            }))
            sys.exit(1)
        
        # Initialize recommender
        recommender = AdvancedCourseRecommender()
        
        # Generate recommendations
        result = recommender.generate_recommendations(
            user_data=user_data,
            users_under_manager=users_under_manager,
            all_courses=all_courses,
            logged_in_user_id=logged_in_user_id
        )
        
        # Output result
        print(json.dumps(result, indent=2))
        
    except Exception as e:
        print(json.dumps({
            "error": str(e),
            "recommendations": [],
            "token_count": 0
        }), file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()