from PIL import Image, ImageEnhance, ImageOps
import sys
import time
import re

# Default ASCII gradients
GRADIENT_DETAILED = " .'`^\",:;Il!i><~+_-?][}{1)(|/tfjrxnuvczXYUJCLQ0OZmwqpdbkhao*#MW&8%B@$"
GRADIENT_BLOCK = " ░▒▓█"
GRADIENT_DEFAULT = GRADIENT_BLOCK

def convert_image(
    image_path, 
    width=100, 
    contrast=1.5, 
    brightness=1.2, 
    sharpness=2.0, 
    gradient=GRADIENT_DEFAULT, 
    inversed=False,
    aspect_correction=0.55,
    dither=True,
    color=False  # New parameter for Color support
):
    """
    Converts an image to an ASCII art string.
    """
    try:
        img = Image.open(image_path)
    except Exception as e:
        raise ValueError(f"Unable to open image file {image_path}. Error: {e}")

    # Calculate dimensions
    original_width, original_height = img.size
    
    # If width is None or <= 0, default to 100
    if not width or width <= 0:
        width = 100

    height = int((original_height / original_width) * width * aspect_correction)
    
    # Resize (LANCZOS for quality)
    img = img.resize((width, height), Image.Resampling.LANCZOS)
    
    # Pre-processing
    if contrast != 1.0:
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(contrast)
    if brightness != 1.0:
        enhancer = ImageEnhance.Brightness(img)
        img = enhancer.enhance(brightness)
    if sharpness != 1.0:
        enhancer = ImageEnhance.Sharpness(img)
        img = enhancer.enhance(sharpness)
    
    # Ensure RGB
    img = img.convert("RGB")
    
    # Handle Gradient
    chars = gradient
    if inversed:
         chars = gradient[::-1]
    num_chars = len(chars)

    # ----------------------------------------------------
    # Step 1: Character Mapping (Grayscale logic)
    # ----------------------------------------------------
    
    # We always need the mapped characters first.
    # We can reuse the dither logic for character selection
    
    char_rows = [] # Will hold just the characters without color
    
    if dither and num_chars > 1:
        # High Quality Dithering for Character Selection
        quantized = img.quantize(colors=num_chars, method=Image.Quantize.MEDIANCUT, kmeans=0, dither=Image.Dither.FLOYDSTEINBERG)
        palette = quantized.getpalette()[:num_chars*3]
        
        # Map palette to brightness
        color_brightness = []
        for i in range(0, len(palette), 3):
            r, g, b = palette[i], palette[i+1], palette[i+2]
            lum = 0.299*r + 0.587*g + 0.114*b
            color_brightness.append((lum, i // 3))
            
        color_brightness.sort(key=lambda x: x[0])
        rank_map = { original_index: rank for rank, (lum, original_index) in enumerate(color_brightness) }
        
        pixels_indices = list(quantized.getdata())
        unique_indices = len(color_brightness)
        full_gradient_len = len(chars)
        
        for i in range(0, len(pixels_indices), width):
            row_indices = pixels_indices[i : i + width]
            row_chars = []
            for idx in row_indices:
                rank = rank_map.get(idx, 0)
                if unique_indices > 1:
                    char_idx = int(rank / (unique_indices - 1) * (full_gradient_len - 1))
                else:
                    char_idx = 0
                row_chars.append(chars[char_idx])
            char_rows.append(row_chars)
    else:
        # Simple Grayscale (no dither)
        gray_img = img.convert("L")
        gray_pixels = list(gray_img.getdata())
        for i in range(0, len(gray_pixels), width):
            row = gray_pixels[i:i+width]
            row_chars = []
            for p in row:
                index = int(p / 255 * (num_chars - 1))
                index = max(0, min(index, num_chars - 1))
                row_chars.append(chars[index])
            char_rows.append(row_chars)

    # ----------------------------------------------------
    # Step 2: Coloring
    # ----------------------------------------------------
    
    if not color:
        # Just join the characters
        return "\n".join(["".join(row) for row in char_rows])
    else:
        # Map colors from original RGB image
        rgb_pixels = list(img.getdata())
        ascii_rows = []
        
        pixel_idx = 0
        for r_idx, row_chars in enumerate(char_rows):
            row_str = ""
            for c_idx, char in enumerate(row_chars):
                r, g, b = rgb_pixels[pixel_idx]
                pixel_idx += 1
                
                # ANSI TrueColor: \033[38;2;R;G;Bm
                color_code = f"\033[38;2;{r};{g};{b}m"
                reset_code = "\033[0m"
                
                row_str += f"{color_code}{char}{reset_code}"
            
            ascii_rows.append(row_str)
            
        return "\n".join(ascii_rows)

def print_art(ascii_art, delay=0.04, typewriter=False):
    """
    Prints ASCII art to the terminal with optional animation.
    """
    if typewriter:
        # If typewriter effect is simpler with color, we need to split by ANSI codes + char
        # Regex to split preserving delimiters: (ANSI code) or (Single Char)
        # Actually simpler: Find all (ANSI Code + Char) OR (Char)
        # But convert_image output format is: [Color][Char][Reset]
        # Or just [Char]
        
        # Regex to match ANSI escape codes: \x1b\[[0-9;]*m
        # We split the string, keeping the delimiters.
        tokens = re.split(r'(\x1b\[[0-9;]*m)', ascii_art)
        
        # tokens will look like: ['', '\x1b[...m', 'Char', '\x1b[0m', 'Char', ...]
        # We want to iterate and print. If it's an ANSI code, just print it (no delay).
        # If it's text, print with delay.
        
        for token in tokens:
            if not token: continue
            sys.stdout.write(token)
            sys.stdout.flush()
            
            # Delay only if it's NOT an escape sequence and NOT a newline (unless we want newline delay)
            if not token.startswith('\x1b'):
                if '\n' in token:
                     # This token might be just a newline or contain one
                     time.sleep(delay)
                else:
                    # Character delay
                    time.sleep(delay / 5) 
    else:
        for line in ascii_art.splitlines():
            print(line)
            sys.stdout.flush() 
            if delay > 0:
                time.sleep(delay)

def display(
    image_path, 
    width=100, 
    contrast=1.5, 
    brightness=1.2, 
    sharpness=2.0, 
    gradient=GRADIENT_DEFAULT, 
    inversed=False,
    aspect_correction=0.55,
    dither=True,
    color=False, # New param
    delay=0.04,
    typewriter=False
):
    """
    High-level function to convert and print an image in one go.
    Recommended for library usage.
    """
    ascii_art = convert_image(
        image_path=image_path,
        width=width,
        contrast=contrast,
        brightness=brightness,
        sharpness=sharpness,
        gradient=gradient,
        inversed=inversed,
        aspect_correction=aspect_correction,
        dither=dither,
        color=color
    )
    print_art(ascii_art, delay=delay, typewriter=typewriter)
