import os
import requests
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
from pathlib import Path
from io import BytesIO
import itertools

class ChessAssetProcessor:
    def __init__(self):
        self.base_dir = Path.cwd()
        self.raw_dir = self.base_dir / 'raw'
        self.training_dir = self.base_dir / 'training'
        
        # URLs for pieces
        self.piece_base_url = 'https://www.chess.com/chess-themes/pieces/neo/300'
        
        # Piece configurations - matches the Swift enum exactly
        self.pieces = {
            'white': ['pawn', 'knight', 'bishop', 'rook', 'queen', 'king'],
            'black': ['pawn', 'knight', 'bishop', 'rook', 'queen', 'king']
        }
        
        # Square colors - used for background variations only
        self.square_colors = {
            'light': '#eeeed2',
            'dark': '#759656',
            'light_highlighted': '#f6f68d',
            'dark_highlighted': '#bdcc49'
        }
        
        # Border configurations
        self.border_configs = [
            {},  # No borders
            {'top': True},
            {'bottom': True},
            {'left': True},
            {'top': True, 'left': True},
            {'bottom': True, 'left': True},
        ]
        
        self._setup_directories()
    
    def _setup_directories(self):
        """Create directory structure for training data"""
        (self.raw_dir / 'pieces').mkdir(parents=True, exist_ok=True)
        self.training_dir.mkdir(exist_ok=True)
        
        # Create directories for each piece type
        for color in ['white', 'black']:
            for piece in self.pieces[color]:
                (self.training_dir / f"{color}_{piece}").mkdir(parents=True, exist_ok=True)
    
    def hex_to_rgb(self, hex_color):
        """Convert hex color to RGB tuple"""
        hex_color = hex_color.lstrip('#')
        return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    
    def create_base_square(self, color, size=100):
        """Create a square with specified color"""
        rgb_color = self.hex_to_rgb(color)
        return Image.new('RGB', (size, size), rgb_color)
    
    def add_borders(self, image, borders, border_color=(0, 0, 0)):
        """Add borders according to configuration"""
        w, h = image.size
        result = image.copy()
        draw = ImageDraw.Draw(result)
        border_size = 5
        
        if borders.get('top'):
            draw.line([(0, 0), (w-1, 0)], fill=border_color, width=border_size)
        if borders.get('bottom'):
            draw.line([(0, h-1), (w-1, h-1)], fill=border_color, width=border_size)
        if borders.get('left'):
            draw.line([(0, 0), (0, h-1)], fill=border_color, width=border_size)
        
        return result
    
    def download_piece(self, color, piece, target_size):
        """Download a specific chess piece"""
        piece_letter = piece[0] if piece != 'knight' else 'n'
        piece_url = f"{self.piece_base_url}/{color[0]}{piece_letter}.png"
        
        response = requests.get(piece_url)
        if response.status_code == 200:
            piece_img = Image.open(BytesIO(response.content)).convert('RGBA')
            return piece_img.resize((target_size, target_size), Image.Resampling.LANCZOS)
        return None
    
    def create_training_data(self):
        """Create comprehensive training dataset"""
        square_size = 100
        
        print("Downloading pieces and creating variations...")
        for color in ['white', 'black']:
            for piece in self.pieces[color]:
                piece_dir = self.training_dir / f"{color}_{piece}"
                print(f"\nProcessing {color} {piece}...")
                
                # Download piece
                piece_img = self.download_piece(color, piece, square_size)
                if piece_img is None:
                    print(f"Failed to download {color} {piece}")
                    continue
                
                # Create variations with different backgrounds
                variation_count = 0
                for bg_name, bg_color in self.square_colors.items():
                    # Create base square
                    base_square = self.create_base_square(bg_color, square_size)
                    
                    # Add border variations
                    for border_config in self.border_configs:
                        # Add borders
                        bordered = self.add_borders(base_square, border_config)
                        
                        # Convert to RGBA for composition
                        bordered_rgba = bordered.convert('RGBA')
                        
                        # Combine with piece
                        combined = Image.alpha_composite(bordered_rgba, piece_img)
                        
                        # Generate filename
                        border_desc = '_'.join(k for k,v in border_config.items() if v)
                        filename = f"{bg_name}_{border_desc}_{variation_count}.png" if border_desc else f"{bg_name}_{variation_count}.png"
                        
                        # Save image
                        combined.save(piece_dir / filename)
                        variation_count += 1
                
                print(f"Created {variation_count} variations for {color} {piece}")

def main():
    processor = ChessAssetProcessor()
    processor.create_training_data()

if __name__ == '__main__':
    main()