import os
import sys
import math
import logging
import concurrent.futures

import numpy as np
from tqdm import tqdm
from PIL import Image, ImageOps

from koba.core import charsets, unify

def process_block(args):
    block, characters, engine, save_chars = args
    return unify.get_character(block, characters, save_chars, engine)

def calculate_block_sizes(width, height, char_aspect, scale):
    terminal_width = os.get_terminal_size().columns
    chars_width = terminal_width
    
    min_block_width = 10 / char_aspect
    min_block_height = 10

    scale = min(scale, 1.0)
    
    max_chars_width_by_block_width = int(width // min_block_width)
    max_chars_width_by_block_height = int(height // min_block_height * char_aspect)

    max_chars_width = min(max_chars_width_by_block_width, max_chars_width_by_block_height)

    chars_width = min(int(terminal_width * scale), terminal_width, max_chars_width)

    chars_width = max(chars_width, 1)

    chars_height = math.ceil((height * chars_width / width) / char_aspect)
    
    block_widths = [width // chars_width] * chars_width
    for i in range(width % chars_width):
        block_widths[i] += 1

    block_heights = [height // chars_height] * chars_height
    for i in range(height % chars_height):
        block_heights[i] += 1
        
    return block_widths, block_heights, chars_width

def process(img, char_aspect, scale, engine, color, invert, stretch_contrast, save_blocks, start_char, end_char, save_chars, font, single_threaded, show_progress=False):
    if color:
        img_color = img.convert("RGB")
        img_arr_color = np.array(img_color)
    
    img = img.convert("L")
    if invert:
        img = ImageOps.invert(img)
    if stretch_contrast:
        img = ImageOps.autocontrast(img)
    
    img_arr = np.array(img)
    height, width = img_arr.shape[:2]
    logging.debug(f"Image is {width}x{height} pixels.")
    
    block_widths, block_heights, chars_width = calculate_block_sizes(width, height, char_aspect, scale)

    logging.debug(f"Image will be {chars_width}x{len(block_heights)} chars.")
    logging.debug(f"Block widths: {block_widths[:5]}... (total {len(block_widths)})")
    logging.debug(f"Block heights: {block_heights[:5]}... (total {len(block_heights)})")

    if min(block_heights + block_widths) <= 7 and str(engine).lower() == "ssim":
        logging.critical("Image blocks are too small for SSIM. Please use another engine.")
        sys.exit(1)

    # splitting into blocks with variable sizes
    block_colors = []
    blocks = []
    y = 0
    for bh in block_heights:
        x = 0
        for bw in block_widths:
            block = img_arr[y:y+bh, x:x+bw]                
            blocks.append(block)
            x += bw
        y += bh
    
    if color:
        y = 0
        for bh in block_heights:
            x = 0
            for bw in block_widths:
                block = img_arr_color[y:y+bh, x:x+bw]
                block_color = block.mean(axis=(0, 1)).astype(int)            
                block_colors.append(block_color)
                x += bw
            y += bh
    
    if color:
        if len(block_colors) != len(blocks):
            logging.critical("Block colors don't match the blocks. Falling back to grayscale.")
            color = False
        else:
            logging.debug(f"Block colors: {block_colors[:3]} ...")
    
    logging.debug(f"{len(blocks)} blocks created.")
    
    if save_blocks:
        logging.info("Saving image blocks...")
        blocks_dir = "blocks"
        os.makedirs(blocks_dir, exist_ok=True)
        for idx, block in enumerate(blocks):
            img_block = Image.fromarray(block)
            img_block.save(os.path.join(blocks_dir, f"{idx:04d}.png"))
    
    print_chars = ""
    
    if start_char == end_char:
        all_chars = chr(start_char) * len(blocks)
    else:
        characters = charsets.get_range(start_char, end_char)
        all_chars = ""

        # prepare arguments
        unique_blocks = {block.tobytes(): block for block in blocks}
        unique_block_args = [(block, characters, engine.lower(), save_chars) for block in unique_blocks.values()]

        if font:
            unify.font_path = font

        block_results = {}
        if not single_threaded:
            with concurrent.futures.ProcessPoolExecutor() as executor:
                future_to_block = {
                    executor.submit(unify.get_character, *arg): arg[0] for arg in unique_block_args
                }
                try:
                    for future in tqdm(
                        concurrent.futures.as_completed(future_to_block),
                        total=len(future_to_block),
                        desc="Processing unique blocks",
                        disable=not show_progress
                    ):
                        block = future_to_block[future]
                        block_results[block.tobytes()] = future.result()
                except ValueError as e:
                    logging.critical(str(e))
                    sys.exit(1)
        else:
            for args in tqdm(
                unique_block_args,
                desc="Processing unique blocks (single-threaded)",
                disable=not show_progress
            ):
                block, characters, engine, save_chars = args
                result = unify.get_character(block, characters, save_chars, engine)
                block_results[block.tobytes()] = result

        results = [block_results[block.tobytes()] for block in blocks]
        
        all_chars = "".join(results)
    
    lines = [all_chars[i:i+chars_width] for i in range(0, len(all_chars), chars_width)]
    flat_chars = "".join(lines)
    logging.debug(f"Char count: {len(flat_chars)}")
    logging.debug(f"Colored blocks: {len(block_colors)}")

    color_idx = 0
    for c in '\n'.join(lines):
        if c == '\n':
            print_chars += '\n'
        elif color_idx < len(block_colors):
            r, g, b = block_colors[color_idx]
            print_chars += f"\033[38;2;{r};{g};{b}m{c}\033[0m"
            color_idx += 1
        else:
            print_chars += c
            
    return print_chars
