"""
King Safety Analysis Tool
Evaluates king safety and exposure to threats.

Returns structured analysis of:
- Pawn shield quality around the king
- Number of attackers vs defenders near the king
- King exposure (open files, diagonals)
- Escape squares availability
- Castling rights and status
"""

import chess
from typing import Dict, List, Set


def analyze_king_safety(fen: str) -> Dict:
    """
    Analyze king safety for both sides.
    
    Args:
        fen: Position in FEN notation
    
    Returns:
        Dictionary with:
            - score: Safety differential (positive = white king safer)
            - factors: List of contributing factors (pawn shield, attackers, etc.)
            - evidence: Human-readable explanations
            - details: Detailed breakdown by side
    
    Example:
        >>> result = analyze_king_safety("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1")
        >>> result['evidence']
        ['Both kings in center, not yet castled', ...]
    """
    board = chess.Board(fen)
    
    # Analyze both kings
    white_safety = _analyze_king_safety_for_side(board, chess.WHITE)
    black_safety = _analyze_king_safety_for_side(board, chess.BLACK)
    
    # Calculate safety differential
    score = white_safety['score'] - black_safety['score']
    
    factors = []
    evidence = []
    
    # Castling status
    _analyze_castling_status(board, white_safety, black_safety, evidence, factors)
    
    # Pawn shield comparison
    white_shield = white_safety['pawn_shield']
    black_shield = black_safety['pawn_shield']
    
    if white_shield > black_shield:
        factors.append({
            'type': 'pawn_shield_advantage',
            'side': 'white',
            'value': white_shield - black_shield
        })
        evidence.append(f"White king has better pawn shield ({white_shield} vs {black_shield} pawns)")
    elif black_shield > white_shield:
        factors.append({
            'type': 'pawn_shield_advantage',
            'side': 'black',
            'value': black_shield - white_shield
        })
        evidence.append(f"Black king has better pawn shield ({black_shield} vs {white_shield} pawns)")
    
    # Attacker count comparison
    white_attackers = white_safety['nearby_attackers']
    black_attackers = black_safety['nearby_attackers']
    
    if white_attackers > 0:
        factors.append({
            'type': 'king_under_attack',
            'side': 'white',
            'attackers': white_attackers
        })
        evidence.append(f"White king faces {white_attackers} nearby attacking piece(s)")
    
    if black_attackers > 0:
        factors.append({
            'type': 'king_under_attack',
            'side': 'black',
            'attackers': black_attackers
        })
        evidence.append(f"Black king faces {black_attackers} nearby attacking piece(s)")
    
    # Escape squares
    white_escapes = white_safety['escape_squares']
    black_escapes = black_safety['escape_squares']
    
    if white_escapes < 2:
        factors.append({
            'type': 'limited_escape_squares',
            'side': 'white',
            'count': white_escapes
        })
        evidence.append(f"White king has limited escape squares ({white_escapes})")
    
    if black_escapes < 2:
        factors.append({
            'type': 'limited_escape_squares',
            'side': 'black',
            'count': black_escapes
        })
        evidence.append(f"Black king has limited escape squares ({black_escapes})")
    
    # Open files near king
    _analyze_open_files_near_king(board, white_safety, black_safety, evidence, factors)
    
    # Check status
    if board.is_check():
        checking_side = "Black" if board.turn == chess.WHITE else "White"
        factors.append({
            'type': 'king_in_check',
            'side': 'white' if board.turn == chess.WHITE else 'black'
        })
        evidence.append(f"{checking_side} has the king in check!")
    
    return {
        'score': score,
        'factors': factors,
        'evidence': evidence,
        'details': {
            'white': white_safety,
            'black': black_safety
        }
    }


def _analyze_king_safety_for_side(board: chess.Board, color: chess.Color) -> Dict:
    """
    Analyze king safety for one side.
    
    Args:
        board: Chess board
        color: Color to analyze
    
    Returns:
        Dictionary with safety metrics
    """
    # Find king position
    king_square = board.king(color)
    
    if king_square is None:
        # No king (shouldn't happen in valid position)
        return {
            'king_square': None,
            'pawn_shield': 0,
            'escape_squares': 0,
            'nearby_attackers': 0,
            'open_files': 0,
            'castled': False,
            'score': 0
        }
    
    king_file = chess.square_file(king_square)
    king_rank = chess.square_rank(king_square)
    
    # Determine if castled (heuristic: king on g-file or c-file, not on starting square)
    starting_square = chess.E1 if color == chess.WHITE else chess.E8
    castled = king_square != starting_square and (king_file in [2, 6])  # c or g file
    
    # Count pawn shield
    pawn_shield = _count_pawn_shield(board, king_square, color)
    
    # Count escape squares
    escape_squares = _count_escape_squares(board, king_square, color)
    
    # Count nearby attackers (enemy pieces attacking squares around king)
    nearby_attackers = _count_nearby_attackers(board, king_square, color)
    
    # Count open files near king
    open_files = _count_open_files_near_king(board, king_square, color)
    
    # Calculate safety score
    # Higher score = safer king
    # Pawn shield: +30 points per pawn
    # Castled: +50 points
    # Escape squares: +15 points per square
    # Nearby attackers: -40 points per attacker
    # Open files near king: -25 points per file
    score = (
        pawn_shield * 30 +
        (50 if castled else 0) +
        escape_squares * 15 -
        nearby_attackers * 40 -
        open_files * 25
    )
    
    return {
        'king_square': king_square,
        'king_position': chess.square_name(king_square),
        'pawn_shield': pawn_shield,
        'escape_squares': escape_squares,
        'nearby_attackers': nearby_attackers,
        'open_files': open_files,
        'castled': castled,
        'score': score
    }


def _count_pawn_shield(board: chess.Board, king_square: int, color: chess.Color) -> int:
    """
    Count pawns in front of the king (pawn shield).
    
    Args:
        board: Chess board
        king_square: Square where king is located
        color: King's color
    
    Returns:
        Number of friendly pawns protecting the king
    """
    king_file = chess.square_file(king_square)
    king_rank = chess.square_rank(king_square)
    
    pawn_count = 0
    
    # Check files: king's file and adjacent files
    files_to_check = [king_file - 1, king_file, king_file + 1]
    
    # Determine ranks to check (in front of king)
    if color == chess.WHITE:
        ranks_to_check = [king_rank + 1, king_rank + 2]
    else:
        ranks_to_check = [king_rank - 1, king_rank - 2]
    
    for file in files_to_check:
        if file < 0 or file > 7:
            continue
        for rank in ranks_to_check:
            if rank < 0 or rank > 7:
                continue
            
            square = chess.square(file, rank)
            piece = board.piece_at(square)
            
            if piece and piece.piece_type == chess.PAWN and piece.color == color:
                pawn_count += 1
    
    return pawn_count


def _count_escape_squares(board: chess.Board, king_square: int, color: chess.Color) -> int:
    """
    Count safe escape squares for the king.
    
    Args:
        board: Chess board
        king_square: Square where king is located
        color: King's color
    
    Returns:
        Number of safe squares the king can move to
    """
    escape_count = 0
    
    # All squares around the king
    king_file = chess.square_file(king_square)
    king_rank = chess.square_rank(king_square)
    
    for file_offset in [-1, 0, 1]:
        for rank_offset in [-1, 0, 1]:
            if file_offset == 0 and rank_offset == 0:
                continue  # Skip king's current square
            
            target_file = king_file + file_offset
            target_rank = king_rank + rank_offset
            
            if target_file < 0 or target_file > 7 or target_rank < 0 or target_rank > 7:
                continue
            
            target_square = chess.square(target_file, target_rank)
            
            # Check if square is occupied by friendly piece
            piece = board.piece_at(target_square)
            if piece and piece.color == color:
                continue
            
            # Check if square is attacked by opponent
            if not board.is_attacked_by(not color, target_square):
                escape_count += 1
    
    return escape_count


def _count_nearby_attackers(board: chess.Board, king_square: int, color: chess.Color) -> int:
    """
    Count enemy pieces attacking squares near the king.
    
    Args:
        board: Chess board
        king_square: Square where king is located
        color: King's color
    
    Returns:
        Number of enemy pieces attacking king's vicinity
    """
    king_file = chess.square_file(king_square)
    king_rank = chess.square_rank(king_square)
    
    attackers = set()  # Use set to avoid counting same piece multiple times
    
    # Check all squares in king's vicinity (3x3 area)
    for file_offset in [-1, 0, 1]:
        for rank_offset in [-1, 0, 1]:
            target_file = king_file + file_offset
            target_rank = king_rank + rank_offset
            
            if target_file < 0 or target_file > 7 or target_rank < 0 or target_rank > 7:
                continue
            
            target_square = chess.square(target_file, target_rank)
            
            # Find all enemy pieces attacking this square
            for attacker_square in board.attackers(not color, target_square):
                piece = board.piece_at(attacker_square)
                if piece and piece.piece_type != chess.PAWN:  # Count major/minor pieces, not pawns
                    attackers.add(attacker_square)
    
    return len(attackers)


def _count_open_files_near_king(board: chess.Board, king_square: int, color: chess.Color) -> int:
    """
    Count open files near the king (potential rook invasion routes).
    
    Args:
        board: Chess board
        king_square: Square where king is located
        color: King's color
    
    Returns:
        Number of open files adjacent to king
    """
    king_file = chess.square_file(king_square)
    
    open_file_count = 0
    
    # Check king's file and adjacent files
    for file_offset in [-1, 0, 1]:
        file_to_check = king_file + file_offset
        
        if file_to_check < 0 or file_to_check > 7:
            continue
        
        # Count pawns on this file
        pawn_count = 0
        for rank in range(8):
            square = chess.square(file_to_check, rank)
            piece = board.piece_at(square)
            if piece and piece.piece_type == chess.PAWN:
                pawn_count += 1
        
        # Open file = no pawns
        if pawn_count == 0:
            open_file_count += 1
    
    return open_file_count


def _analyze_castling_status(board: chess.Board, white_safety: Dict, black_safety: Dict, 
                            evidence: List[str], factors: List[Dict]):
    """
    Analyze and report on castling status.
    
    Args:
        board: Chess board
        white_safety: White king safety data
        black_safety: Black king safety data
        evidence: List to append evidence to
        factors: List to append factors to
    """
    # Check castling rights
    white_can_castle_kingside = board.has_kingside_castling_rights(chess.WHITE)
    white_can_castle_queenside = board.has_queenside_castling_rights(chess.WHITE)
    black_can_castle_kingside = board.has_kingside_castling_rights(chess.BLACK)
    black_can_castle_queenside = board.has_queenside_castling_rights(chess.BLACK)
    
    # White castling status
    if white_safety['castled']:
        evidence.append(f"White king has castled (on {white_safety['king_position']})")
        factors.append({
            'type': 'castled',
            'side': 'white',
            'position': white_safety['king_position']
        })
    elif white_can_castle_kingside or white_can_castle_queenside:
        castling_options = []
        if white_can_castle_kingside:
            castling_options.append("kingside")
        if white_can_castle_queenside:
            castling_options.append("queenside")
        evidence.append(f"White can still castle {' or '.join(castling_options)}")
    else:
        evidence.append("White has lost castling rights")
        factors.append({
            'type': 'lost_castling_rights',
            'side': 'white'
        })
    
    # Black castling status
    if black_safety['castled']:
        evidence.append(f"Black king has castled (on {black_safety['king_position']})")
        factors.append({
            'type': 'castled',
            'side': 'black',
            'position': black_safety['king_position']
        })
    elif black_can_castle_kingside or black_can_castle_queenside:
        castling_options = []
        if black_can_castle_kingside:
            castling_options.append("kingside")
        if black_can_castle_queenside:
            castling_options.append("queenside")
        evidence.append(f"Black can still castle {' or '.join(castling_options)}")
    else:
        evidence.append("Black has lost castling rights")
        factors.append({
            'type': 'lost_castling_rights',
            'side': 'black'
        })


def _analyze_open_files_near_king(board: chess.Board, white_safety: Dict, black_safety: Dict,
                                 evidence: List[str], factors: List[Dict]):
    """
    Analyze and report on open files near kings.
    
    Args:
        board: Chess board
        white_safety: White king safety data
        black_safety: Black king safety data
        evidence: List to append evidence to
        factors: List to append factors to
    """
    if white_safety['open_files'] > 0:
        factors.append({
            'type': 'open_files_near_king',
            'side': 'white',
            'count': white_safety['open_files']
        })
        evidence.append(f"White king exposed to {white_safety['open_files']} open file(s)")
    
    if black_safety['open_files'] > 0:
        factors.append({
            'type': 'open_files_near_king',
            'side': 'black',
            'count': black_safety['open_files']
        })
        evidence.append(f"Black king exposed to {black_safety['open_files']} open file(s)")
