ATS Resume Optimizer Agent

"""
Unified Agent Template - Works with OpenAI, Anthropic, and Local WebLLM
Pure Python tools with conditional LangChain wrapping based on provider.

This template enables ONE codebase for all three providers:
- OpenAI: Uses LangChain with ChatOpenAI
- Anthropic: Uses LangChain with ChatAnthropic  
- Local WebLLM: Routes to JavaScript LangChain.js bridge (NO Python LangChain)

Key features:
- Pure Python tool functions (no decorators at definition time)
- Schema extraction via get_tool_schemas() for WebLLM JavaScript bridge
- Conditional LangChain imports only when needed (cloud providers)
- Runtime tool wrapping with tool() function for cloud providers
"""

import json
import inspect
from typing import get_type_hints

# Provider injected from generator context (openai|anthropic|local)

# ============================================================================
# Schema Extraction Helpers (for all providers)
# ============================================================================

def _python_type_to_json_type(python_type):
    """Convert Python type to JSON schema type."""
    type_mapping = {
        'str': 'string',
        'int': 'integer',
        'float': 'number',
        'bool': 'boolean',
        'list': 'array',
        'dict': 'object',
    }
    type_name = python_type.__name__ if hasattr(python_type, '__name__') else str(python_type)
    return type_mapping.get(type_name, 'string')


def _extract_function_schema(func):
    """
    Extract OpenAI function schema from a Python function.
    Works for all providers - no LangChain dependency.
    
    Extracts:
    - Function name
    - Description from docstring
    - Parameters with types and descriptions
    - Required vs optional parameters
    
    Args:
        func: Python function with type hints and docstring
    
    Returns:
        dict: OpenAI function schema format
    """
    sig = inspect.signature(func)
    
    # Get type hints
    try:
        hints = get_type_hints(func)
    except:
        hints = {}
    
    # Parse docstring
    doc = inspect.getdoc(func) or ""
    description = doc.split('\n\n')[0] if doc else f"Execute {func.__name__}"
    
    # Build parameters schema
    parameters = {
        "type": "object",
        "properties": {},
        "required": []
    }
    
    # Extract parameter descriptions from docstring Args section
    param_descriptions = {}
    if "Args:" in doc:
        args_section = doc.split("Args:")[1].split("Returns:")[0] if "Returns:" in doc else doc.split("Args:")[1]
        for line in args_section.split('\n'):
            line = line.strip()
            if ':' in line:
                param_name = line.split(':')[0].strip()
                param_desc = line.split(':', 1)[1].strip()
                param_descriptions[param_name] = param_desc
    
    # Process each parameter
    for param_name, param in sig.parameters.items():
        if param_name in ['self', 'cls']:
            continue
        
        param_type = hints.get(param_name, param.annotation)
        if param_type == inspect.Parameter.empty:
            param_type = str
        
        json_type = _python_type_to_json_type(param_type)
        
        prop_schema = {
            "type": json_type,
            "description": param_descriptions.get(param_name, f"The {param_name} parameter")
        }
        
        parameters["properties"][param_name] = prop_schema
        
        # Mark as required if no default value
        if param.default == inspect.Parameter.empty:
            parameters["required"].append(param_name)
    
    return {
        "name": func.__name__,
        "description": description,
        "parameters": parameters
    }

# Resume Optimizer AI Agent
# Analyzes resumes against job descriptions and provides optimization recommendations

import json
import re
from typing import Dict, List, Any

# Global state
api_key = ""

# Import LangChain
try:
    from langchain_core.messages import HumanMessage, SystemMessage
    from langchain_openai import ChatOpenAI
    langchain_available = True
    print("[OK] LangChain imports successful")
except ImportError as e:
    langchain_available = False
    print(f"[WARNING] LangChain not available: {e}")

def extract_keywords_from_job(job_description: str) -> Dict[str, List[str]]:
    """Extract required skills and keywords from job description"""
    
    # Common keyword patterns
    skill_patterns = [
        r'(?i)(?:experience with|proficiency in|knowledge of|familiar with|skilled in)\s+([\w\s,/+#-]+)',
        r'(?i)(?:required|must have|should have):\s*([\w\s,/+#-]+)',
        r'(?i)(?:skills?|technologies?|tools?|languages?):\s*([\w\s,/+#-]+)'
    ]
    
    keywords = {
        'critical': [],
        'high': [],
        'medium': [],
        'nice_to_have': []
    }
    
    # Extract from "Requirements" section
    requirements_section = re.search(r'(?i)requirements?:(.*?)(?:nice to have|$)', job_description, re.DOTALL)
    if requirements_section:
        req_text = requirements_section.group(1)
        # Extract bullet points
        bullets = re.findall(r'[•●▪️-]\s*(.+)', req_text)
        
        for bullet in bullets:
            # Extract skills/technologies
            words = re.findall(r'\b[A-Z][\w+#.-]+\b', bullet)
            if words:
                if any(term in bullet.lower() for term in ['required', 'must', 'essential']):
                    keywords['critical'].extend(words)
                elif any(term in bullet.lower() for term in ['experience', 'years', 'proven']):
                    keywords['high'].extend(words)
                else:
                    keywords['medium'].extend(words)
    
    # Extract from "Nice to have" section
    nice_to_have = re.search(r'(?i)nice to have:(.*?)$', job_description, re.DOTALL)
    if nice_to_have:
        nice_text = nice_to_have.group(1)
        words = re.findall(r'\b[A-Z][\w+#.-]+\b', nice_text)
        keywords['nice_to_have'].extend(words)
    
    # Remove duplicates
    for category in keywords:
        keywords[category] = list(set(keywords[category]))
    
    return keywords

def extract_resume_keywords(resume_text: str) -> List[str]:
    """Extract technical keywords from resume"""
    
    # Common technical terms pattern
    tech_keywords = re.findall(r'\b[A-Z][\w+#.-]+\b', resume_text)
    
    # Filter out common words
    common_words = {'The', 'And', 'For', 'With', 'This', 'That', 'From', 'Have', 'Been', 'Were', 'Will'}
    filtered = [kw for kw in tech_keywords if kw not in common_words]
    
    return list(set(filtered))

def calculate_keyword_density(resume_text: str, keywords: List[str]) -> float:
    """Calculate keyword density percentage"""
    
    total_words = len(resume_text.split())
    keyword_count = sum(resume_text.lower().count(kw.lower()) for kw in keywords)
    
    if total_words == 0:
        return 0.0
    
    density = (keyword_count / total_words) * 100
    return round(density, 2)

def match_keywords(resume_keywords: List[str], job_keywords: Dict[str, List[str]]) -> Dict[str, List[str]]:
    """Match resume keywords against job requirements"""
    
    resume_lower = [kw.lower() for kw in resume_keywords]
    all_job_keywords = []
    for category in job_keywords.values():
        all_job_keywords.extend([kw.lower() for kw in category])
    
    matched = [kw for kw in resume_keywords if kw.lower() in all_job_keywords]
    missing = [kw for kw in all_job_keywords if kw not in resume_lower]
    
    return {
        'matched': matched,
        'missing': list(set(missing))
    }

def calculate_ats_score(resume_text: str, job_description: str, matched_keywords: Dict[str, List[str]]) -> Dict[str, Any]:
    """Calculate ATS compatibility score"""
    
    # Extract keywords
    job_keywords = extract_keywords_from_job(job_description)
    resume_keywords = extract_resume_keywords(resume_text)
    
    # Count total required keywords
    total_critical = len(job_keywords['critical'])
    total_high = len(job_keywords['high'])
    
    # Count matched keywords
    matched_critical = sum(1 for kw in job_keywords['critical'] if kw.lower() in [k.lower() for k in matched_keywords['matched']])
    matched_high = sum(1 for kw in job_keywords['high'] if kw.lower() in [k.lower() for k in matched_keywords['matched']])
    
    # Calculate keyword score (0-100)
    if total_critical + total_high > 0:
        keyword_score = ((matched_critical * 2 + matched_high) / (total_critical * 2 + total_high)) * 100
    else:
        keyword_score = 50
    
    # Format score (simple heuristics)
    format_score = 95  # Assume good formatting
    if len(resume_text.split('\n')) < 10:
        format_score -= 20
    
    # Content score (check for quantifiable achievements)
    achievement_count = len(re.findall(r'\d+%|\d+\+|\$\d+|\d+x', resume_text))
    content_score = min(100, 60 + (achievement_count * 5))
    
    # Experience score
    years_mentioned = len(re.findall(r'\d+\+?\s*years?', resume_text, re.IGNORECASE))
    experience_score = min(100, 70 + (years_mentioned * 10))
    
    # Overall score (weighted average)
    overall_before = int(
        keyword_score * 0.4 +
        format_score * 0.2 +
        content_score * 0.2 +
        experience_score * 0.2
    )
    
    # Optimized score (assume +25% improvement)
    overall_after = min(100, overall_before + 25)
    
    return {
        'before': overall_before,
        'after': overall_after,
        'breakdown': {
            'Keywords': int(keyword_score),
            'Formatting': format_score,
            'Content': content_score,
            'Experience': experience_score
        }
    }

def optimize_bullet_points(text: str) -> List[Dict[str, str]]:
    """Identify weak bullet points and suggest improvements"""
    
    weak_verbs = ['worked', 'did', 'was', 'responsible for', 'helped', 'assisted']
    strong_verbs = ['architected', 'engineered', 'led', 'developed', 'implemented', 'designed', 'optimized']
    
    improvements = []
    
    # Find bullet points
    bullets = re.findall(r'[•●▪️-]\s*(.+)', text)
    
    for bullet in bullets:
        bullet_lower = bullet.lower()
        
        # Check for weak verbs
        for weak in weak_verbs:
            if weak in bullet_lower:
                improvements.append({
                    'type': 'weak_verb',
                    'original': bullet,
                    'issue': f'Weak action verb: "{weak}"',
                    'suggestion': f'Use stronger verbs like: {", ".join(strong_verbs[:3])}'
                })
                break
        
        # Check for lack of quantification
        if not re.search(r'\d+', bullet):
            improvements.append({
                'type': 'no_metrics',
                'original': bullet,
                'issue': 'Missing quantifiable results',
                'suggestion': 'Add numbers, percentages, or specific outcomes'
            })
    
    return improvements

def optimize_resume(resume_text: str, job_description: str, industry: str, experience_level: str) -> str:
    """Main optimization function"""
    
    print(f"[OPTIMIZE] Starting optimization for {industry} - {experience_level}")
    
    # Extract keywords
    job_keywords = extract_keywords_from_job(job_description)
    resume_keywords = extract_resume_keywords(resume_text)
    
    print(f"[OPTIMIZE] Found {len(resume_keywords)} keywords in resume")
    print(f"[OPTIMIZE] Found {sum(len(v) for v in job_keywords.values())} keywords in job description")
    
    # Match keywords
    matched = match_keywords(resume_keywords, job_keywords)
    
    # Calculate scores
    ats_score = calculate_ats_score(resume_text, job_description, matched)
    
    # Keyword density
    all_job_kw = []
    for v in job_keywords.values():
        all_job_kw.extend(v)
    density = calculate_keyword_density(resume_text, all_job_kw)
    
    # Identify improvements
    bullet_improvements = optimize_bullet_points(resume_text)
    
    # Generate recommendations
    recommendations = []
    
    # Critical: Missing keywords
    if len(matched['missing']) > 0:
        recommendations.append({
            'type': 'critical',
            'title': 'Missing Required Keywords',
            'text': f"Add these keywords from job description: {', '.join(matched['missing'][:5])}"
        })
    
    # Warnings: Bullet point improvements
    for improvement in bullet_improvements[:3]:
        recommendations.append({
            'type': 'warning',
            'title': improvement['issue'],
            'text': f"Original: \"{improvement['original']}\" - {improvement['suggestion']}"
        })
    
    # Pro tips
    recommendations.append({
        'type': 'info',
        'title': '💡 Pro Tip',
        'text': 'Start each bullet point with a strong action verb and include quantifiable results'
    })
    
    # IMPORTANT (Local/WebLLM mode):
    # The in-browser LLM is responsible for generating the rewritten resume.
    # This tool returns analysis + recommendations that the LLM uses.
    print("[OPTIMIZE] Returning analysis for WebLLM to synthesize optimized resume")
    optimized_text = resume_text
    
    # Prepare result
    result = {
        'optimized_resume': optimized_text,
        'ats_score': ats_score,
        'keyword_analysis': {
            'matched': matched['matched'],
            'missing': matched['missing'],
            'total': len(matched['matched']) + len(matched['missing']),
            'density': density
        },
        'recommendations': recommendations
    }
    
    return json.dumps(result)

print("[INIT] Resume Optimizer ready")
print("[INIT] Upload your resume and paste a job description to begin")

# ============================================================================
# Tool Schema Export (for WebLLM JavaScript bridge)
# ============================================================================

def get_tool_schemas() -> str:
    """
    Export all tool function schemas in OpenAI format.
    Used by WebLLM JavaScript bridge to discover available tools.
    
    This function is called by PyodideToolBridge in JavaScript to:
    1. Discover what tools are available
    2. Get their schemas for LLM binding
    3. Enable function calling in WebLLM
    
    Returns:
        str: JSON string with OpenAI-format tool schemas
    """
    # List all your tool functions here
    # Example: tool_functions = [example_tool, another_tool]
    tool_functions = [
        # TODO: Add your tool functions here
        # For Data Analysis Agent:
        # load_csv_data, get_data_summary, get_column_info, get_value_counts, create_chart, get_correlation_analysis
    ]
    
    # Auto-discovery: if tool_functions is empty, try to find functions defined in this module
    # that are not private (start with _) and not imported
    if not tool_functions:
        import inspect
        import sys
        current_module = sys.modules[__name__]
        for name, obj in inspect.getmembers(current_module):
            if inspect.isfunction(obj) and not name.startswith('_'):
                # Filter out imported functions and infrastructure functions
                if obj.__module__ == __name__ and name not in ['get_tool_schemas', 'process_user_query', 'process_user_query_webllm']:
                    tool_functions.append(obj)
    
    schemas = []
    for func in tool_functions:
        try:
            function_schema = _extract_function_schema(func)
            openai_schema = {
                "type": "function",
                "function": function_schema
            }
            schemas.append(openai_schema)
        except Exception as e:
            print(f"Warning: Failed to extract schema for {func.__name__}: {e}")
    
    return json.dumps(schemas, indent=2)


# ============================================================================
# Unified Query Processing
# ============================================================================

async def process_user_query(query: str) -> str:
    """
    Unified query processor - works for ALL providers (OpenAI, Anthropic, Local).
    
    Routes to appropriate backend based on PROVIDER global variable:
    - 'local': Routes to JavaScript WebLLM agent (NO Python LangChain)
    - 'openai': Uses Python LangChain with ChatOpenAI
    - 'anthropic': Uses Python LangChain with ChatAnthropic
    
    Args:
        query: User's natural language query
    
    Returns:
        str: Response from the agent
    """
    provider = globals().get('PROVIDER', 'openai')
    
    if provider == 'local':
        # ====================================================================
        # WebLLM Local Mode: Use JavaScript LangChain.js bridge
        # ====================================================================
        # No Python LangChain imports needed here
        # All inference happens in JavaScript with WebLLM + LangChain.js
        # Tools are executed in Python via PyodideToolBridge
        
        try:
            # Route to JavaScript WebLLM agent
            # This function is defined in the HTML template and bridges to JS
            result = await process_user_query_webllm(query)
            return result
        except Exception as e:
            print(f"[WebLLM Error] {str(e)}")
            import traceback
            traceback.print_exc()
            return f"❌ Error using WebLLM: {str(e)}"
    
    else:
        # ====================================================================
        # Cloud Providers (OpenAI/Anthropic): Use Python LangChain
        # ====================================================================
        # Import LangChain components ONLY for cloud providers
        # This keeps the bundle smaller for local mode
        
        try:
            from langchain_core.tools import tool
            from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
            
            # Get API key from globals (set by HTML template)
            api_key = globals().get('current_api_key', '')
            if not api_key:
                return "⚠️ Please enter an API key to use AI-powered features."
            
            # Import and initialize provider-specific LLM
            if provider == 'openai':
                from langchain_openai import ChatOpenAI
                llm = ChatOpenAI(
                    model="gpt-3.5-turbo",
                    api_key=api_key,
                    temperature=0.7
                )
            elif provider == 'anthropic':
                from langchain_anthropic import ChatAnthropic
                llm = ChatAnthropic(
                    model="gpt-3.5-turbo",
                    api_key=api_key,
                    temperature=0.7
                )
            else:
                return f"❌ Unknown provider: {provider}"
            
            # Get tool functions list using auto-discovery (same as get_tool_schemas)
            tool_functions = []
            
            # Auto-discover tool functions from current module
            import sys
            current_module = sys.modules[__name__]
            for name, obj in inspect.getmembers(current_module):
                if inspect.isfunction(obj) and not name.startswith('_'):
                    # Filter out imported functions and infrastructure functions
                    if obj.__module__ == __name__ and name not in ['get_tool_schemas', 'process_user_query', 'process_user_query_webllm', '_python_type_to_json_type', '_extract_function_schema']:
                        tool_functions.append(obj)
            
            if not tool_functions:
                # No tools defined - simple conversation mode
                response = llm.invoke([HumanMessage(content=query)])
                return response.content
            
            # Wrap tools with @tool decorator at runtime
            # This is the key: tool() is called as a FUNCTION, not decorator
            tools = [tool(func) for func in tool_functions]
            llm_with_tools = llm.bind_tools(tools)
            
            # Tool calling loop with message history
            messages = [HumanMessage(content=query)]
            
            max_iterations = 3  # Prevent infinite loops
            for iteration in range(max_iterations):
                response = llm_with_tools.invoke(messages)
                messages.append(response)
                
                # Check if LLM made any tool calls
                if not response.tool_calls:
                    break  # No more tools to call, we're done
                
                # Execute each tool call
                for tool_call in response.tool_calls:
                    tool_name = tool_call['name']
                    tool_args = tool_call['args']
                    
                    # Build tool name -> function mapping
                    tool_map = {func.__name__: func for func in tool_functions}
                    
                    if tool_name in tool_map:
                        try:
                            # Execute the tool
                            result = tool_map[tool_name](**tool_args)
                            messages.append(ToolMessage(
                                content=str(result),
                                tool_call_id=tool_call['id']
                            ))
                        except Exception as e:
                            # Tool execution failed
                            messages.append(ToolMessage(
                                content=f"❌ Error executing {tool_name}: {str(e)}",
                                tool_call_id=tool_call['id']
                            ))
                    else:
                        # Unknown tool requested
                        messages.append(ToolMessage(
                            content=f"❌ Unknown tool: {tool_name}",
                            tool_call_id=tool_call['id']
                        ))
            
            # Return final response content
            final_message = messages[-1]
            if hasattr(final_message, 'content'):
                return final_message.content
            else:
                return str(final_message)
                
        except Exception as e:
            import traceback
            traceback.print_exc()
            return f"❌ Error processing query: {str(e)}"


# ============================================================================
# Initialization Message
# ============================================================================

print("✅ Unified agent initialized (provider: {})".format(globals().get('PROVIDER', 'openai')))