from typing import List, Dict, Any, Optional
from env.models import Action

def grade_detect(issues_pred: List[str], expected: List[str]) -> float:
    """
    F1 score for issue detection (normalize lowercase, token overlap).
    """
    if not issues_pred or not expected:
        return 0.0
    
    # Normalize to lowercase and extract tokens
    pred_tokens = set(' '.join(issues_pred).lower().split())
    exp_tokens = set(' '.join(expected).lower().split())
    
    # Calculate precision and recall
    if not pred_tokens:
        return 0.0
    
    intersection = pred_tokens & exp_tokens
    precision = len(intersection) / len(pred_tokens)
    recall = len(intersection) / len(exp_tokens) if exp_tokens else 0
    
    # F1 score
    if precision + recall == 0:
        return 0.0
    return 2 * (precision * recall) / (precision + recall)

def grade_risk(pred: str, expected: str) -> float:
    """
    1.0 exact, 0.5 adjacent, 0.0 mismatch (ordinal {'low':0,'medium':1,'high':2}).
    """
    if not pred or not expected:
        return 0.0
    
    risk_order = {'low': 0, 'medium': 1, 'high': 2}
    
    try:
        pred_val = risk_order[pred.lower()]
        exp_val = risk_order[expected.lower()]
        
        if pred_val == exp_val:
            return 1.0
        elif abs(pred_val - exp_val) == 1:
            return 0.5
        else:
            return 0.0
    except KeyError:
        return 0.0

def levenshtein_distance(s1: str, s2: str) -> int:
    """Calculate Levenshtein distance between two strings."""
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)
    
    if len(s2) == 0:
        return len(s1)
    
    previous_row = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row
    
    return previous_row[-1]

def grade_fix(fix: str, canonical: str, orig: str) -> float:
    """
    Deterministic fix grading with bounded components and normalized output.
    Returns score in [0.0, 1.0].
    
    Components:
    - syntax_ok: 1.0 if syntax is valid, 0.0 otherwise
    - lev_norm: Normalized Levenshtein distance (0.0 to 1.0)
    - bad_pattern_absent: 0.1 if no bad patterns, 0.0 otherwise
    
    Final score: syntax_ok - lev_norm + bad_pattern_absent, clamped to [0.0, 1.0]
    """
    if not fix:
        return 0.0
    
    # Syntax check (basic) - bounded component
    syntax_ok = 1.0
    try:
        # Check for balanced parentheses and brackets
        brackets = {'(': ')', '[': ']', '{': '}'}
        stack = []
        for char in fix:
            if char in brackets:
                stack.append(brackets[char])
            elif char in brackets.values():
                if not stack or stack.pop() != char:
                    syntax_ok = 0.0
                    break
        if stack:
            syntax_ok = 0.0
    except:
        syntax_ok = 0.0
    
    # Normalized Levenshtein distance - bounded component [0.0, 1.0]
    max_len = max(len(fix), len(canonical))
    if max_len == 0:
        lev_norm = 0.0
    else:
        distance = levenshtein_distance(fix.lower(), canonical.lower())
        lev_norm = min(distance / max_len, 1.0)  # Ensure bounded to [0.0, 1.0]
    
    # Check for bad patterns (simple heuristic) - bounded component [0.0, 0.1]
    bad_patterns = ['pass', 'TODO', 'FIXME', 'XXX', 'print(']
    bad_pattern_absent = 0.1
    for pattern in bad_patterns:
        if pattern in fix:
            bad_pattern_absent = 0.0
            break
    
    # Calculate raw score and clamp to [0.0, 1.0]
    raw_score = syntax_ok - lev_norm + bad_pattern_absent
    return max(0.0, min(1.0, raw_score))

def grade_step(action: Action, expected: Dict[str, Any], original_code: str = "") -> Dict[str, Any]:
    """
    Grade a single step with deterministic partial-credit scoring.
    Returns score (0-1) and feedback.
    """
    if action.action_type == "detect":
        score = grade_detect(action.detected_issues, expected["issues"])
        feedback = f"Issue detection score: {score:.2f}"
        if score < 0.5:
            feedback += f" (expected issues: {expected['issues']})"
        
    elif action.action_type == "risk":
        score = grade_risk(action.risk_level, expected["risk"])
        feedback = f"Risk assessment score: {score:.2f}"
        if score < 1.0:
            feedback += f" (expected risk: {expected['risk']})"
        
    elif action.action_type == "fix":
        score = grade_fix(action.suggested_fix, expected["fix_contains"], original_code)
        feedback = f"Fix suggestion score: {score:.2f}"
        if score < 0.7:
            feedback += f" (expected to contain: {expected['fix_contains']})"
        
    else:
        score = 0.0
        feedback = f"Unknown action type: {action.action_type}"
    
    return {
        "score": score,
        "feedback": feedback
    }

def map_score_to_reward(score: float, action_type: str) -> float:
    """
    Map grader score (0-1) to reward values with appropriate weighting.
    """
    # Base weights for each action type
    weights = {
        "detect": 0.3,
        "risk": 0.2,
        "fix": 0.5
    }
    
    base_weight = weights.get(action_type, 0.1)
    
    # Apply score to weight, with small penalty for very low scores
    if score >= 0.8:
        return base_weight
    elif score >= 0.6:
        return base_weight * 0.8
    elif score >= 0.4:
        return base_weight * 0.6
    elif score >= 0.2:
        return base_weight * 0.4
    else:
        return -base_weight * 0.3  # Small penalty for very poor performance

# Legacy function for backward compatibility
def grade_action(action, expected):
    """Legacy grading function - use grade_step instead."""
    result = grade_step(action, expected)
    return result["score"], result["feedback"]