import argparse
import io
import os
import sys
from typing import List, Optional, Tuple, Union

import cv2  # pyright: ignore[reportMissingImports]
import numpy as np
from PIL import Image, ImageGrab

# --- 类型别名定义，保持代码清晰 ---
RgbColor = Tuple[int, int, int]
ColorInput = Union[RgbColor, str]
ImageSource = Union[str, os.PathLike, io.BytesIO]
Coordinate = Tuple[float, float]


def get_image_from_clipboard() -> Optional[io.BytesIO]:
    """从系统剪贴板获取图像，并将其作为内存中的二进制缓冲区返回。"""
    try:
        clipboard_image = ImageGrab.grabclipboard()
        if not isinstance(clipboard_image, Image.Image):
            return None
        buffer = io.BytesIO()
        clipboard_image.save(buffer, format='PNG')
        buffer.seek(0)
        return buffer
    except Exception:
        return None


def _parse_color(color: ColorInput) -> RgbColor:
    """
    内部辅助函数：将灵活的颜色输入（HEX 字符串或 RGB 元组）解析为经过验证的 RGB 元组。
    """
    if isinstance(color, str):
        hex_val = color.lstrip('#')
        if len(hex_val) == 6:
            try:
                # 将 6 位 HEX 字符串转换为 (R, G, B) 元组
                return (
                    int(hex_val[0:2], 16),
                    int(hex_val[2:4], 16),
                    int(hex_val[4:6], 16),
                )
            except ValueError:
                raise ValueError(f'无效的 HEX 颜色字符串: {color}')
        raise ValueError(f'无效的 HEX 颜色格式: {color}。必须是 6 位十六进制代码。')

    if isinstance(color, tuple) and len(color) == 3:
        if all(isinstance(val, int) and 0 <= val <= 255 for val in color):
            return color
        raise ValueError(f'无效的 RGB 值: {color}。每个值都必须是 0 到 255 之间的整数。')

    raise TypeError(f'不支持的颜色类型: {type(color)}。必须是 RGB 元组或 HEX 字符串。')


def filter_image(
    image: np.ndarray,
    target_color: ColorInput,
    color_tolerance: int = 15,
    return_mask: bool = False,
    sv_lower_bound: int | Tuple[int, int] = 40,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """
    根据指定颜色过滤图像，保留特定颜色的部分。

    基于 HSV 色彩空间进行颜色匹配，相比 RGB 空间更适合颜色分割任务。
    使用色相容忍度和饱和度/亮度下界来确保匹配质量。

    Parameters
    ----------
    - `image` : `np.ndarray`
        - 输入的 RGB 图像数组，shape 为 (H, W, 3)
    - `target_color` : `ColorInput`
        - 目标颜色，可以是 HEX 字符串 (如 '#FF0000') 或 RGB 元组 (如 (255, 0, 0))
    - `color_tolerance` : `int`, optional
        - 颜色匹配的容忍度 (0-179)，数值越大匹配范围越广, by default 15
    - `return_mask` : `bool`, optional
        - 是否同时返回二值掩码, by default False
    - `sv_lower_bound` : `int | Tuple[int, int]`, optional
        - 饱和度和亮度的下界，可以是单个整数或二元组 (s_lower, v_lower)

    Returns
    ----------
    - `np.ndarray` or `Tuple[np.ndarray, np.ndarray]`
        - 如果 return_mask=False: 返回过滤后的图像 (背景为黑色)
        - 如果 return_mask=True: 返回 (过滤后的图像, 二值掩码) 元组
    """
    # 解析目标颜色
    try:
        target_rgb = _parse_color(target_color)
    except (ValueError, TypeError) as e:
        raise ValueError(f'颜色参数错误: {e}')

    # 转换到 HSV 色彩空间
    hsv_image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    target_rgb_array = np.array([[target_rgb]], dtype=np.uint8)
    target_hsv_color = cv2.cvtColor(target_rgb_array, cv2.COLOR_RGB2HSV)[0][0]
    h, s, v = target_hsv_color

    s_low, v_low = (
        sv_lower_bound
        if isinstance(sv_lower_bound, tuple)
        else (sv_lower_bound, sv_lower_bound)
    )
    # 设置 HSV 范围
    lower_bound = np.array(
        [max(0, h - color_tolerance), max(s_low, s - 90), max(v_low, v - 90)]
    )
    upper_bound = np.array([min(179, h + color_tolerance), 255, 255])

    # 生成二值掩码
    binary_mask = cv2.inRange(hsv_image, lower_bound, upper_bound)

    # 应用掩码到原图像
    filtered_image = image.copy()
    filtered_image[binary_mask == 0] = [0, 0, 0]  # 将非目标颜色区域设为黑色

    if return_mask:
        return filtered_image, binary_mask
    return filtered_image


def extract_line_data(
    image_input: ImageSource,
    x_range: Tuple[float, float],
    y_range: Tuple[float, float],
    num_points: Optional[int] = None,
    target_color: Optional[ColorInput] = None,
    color_tolerance: int = 15,
) -> List[Coordinate]:
    """
    从图表图像中提取单条数据线。

    支持两种策略：
    1. 默认：使用自适应阈值，适用于颜色不明确或黑白的图像。
    2. 颜色指定：如果提供了 target_color，则精确提取该颜色的线条。

    Parameters
    ----------
    - `image_input`: 输入的图像文件路径或内存缓冲区。
    - `x_range`, `y_range`: X轴和Y轴的数据范围。
    - `num_points`: (可选) 输出数据的重采样点数。
    - `target_color`: (可选) 目标线条的颜色，可以是 HEX 字符串 (如 '#FF0000') 或 RGB 元组 (如 (255, 0, 0))。
    - `color_tolerance`: (可选) 颜色匹配的容忍度，数值越大范围越广。

    Returns/Raises
    ----------
    - A list of (x, y) coordinates.
    - ValueError if the process fails or input is invalid.
    """
    try:
        image = Image.open(image_input)
    except Exception as e:
        raise ValueError(f'无法读取输入图像。错误: {e}')

    img_np = np.array(image.convert('RGB'))

    # --- 步骤 1: 根据输入选择策略，生成二值图像 (binary_mask) ---
    final_target_rgb = None
    if target_color is not None:
        try:
            # 在函数入口处解析颜色，确保后续代码处理的是标准化的 RGB 元组
            final_target_rgb = _parse_color(target_color)
        except (ValueError, TypeError) as e:
            # 将解析错误重新抛出，以便用户可以修复输入
            raise ValueError(f'颜色参数错误: {e}')

    if final_target_rgb:
        # 策略 A: 基于颜色的分割
        _, binary_mask = filter_image(
            img_np, final_target_rgb, color_tolerance, return_mask=True
        )
    else:
        # 策略 B: 自适应阈值 (原始的、稳定的方法)
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        binary_mask = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 4
        )

    # --- 步骤 2: 定位绘图区域 ---
    non_zero_pixels = cv2.findNonZero(binary_mask)
    if non_zero_pixels is None:
        if final_target_rgb:
            raise ValueError(
                f'对于颜色 {final_target_rgb} (容忍度: {color_tolerance}) 未找到任何像素。'
                '请尝试增大 --color-tolerance 或检查颜色值。'
            )
        else:
            raise ValueError('使用自适应阈值未能检测到任何线条或绘图特征。')

    x_px, y_px, w_px, h_px = cv2.boundingRect(non_zero_pixels)
    plot_area = binary_mask[y_px : y_px + h_px, x_px : x_px + w_px]
    px_height, px_width = plot_area.shape

    # --- 步骤 3, 4, 5: 坐标提取、变换和重采样 (逻辑不变) ---
    pixel_coords: List[Tuple[int, float]] = []
    for x_col in range(px_width):
        y_coords_in_col = np.where(plot_area[:, x_col] > 0)[0]
        if y_coords_in_col.size > 0:
            y_center = np.mean(y_coords_in_col)
            pixel_coords.append((x_col, float(y_center)))

    if not pixel_coords:
        return []

    data_coords: List[Coordinate] = []
    x_min_data, x_max_data = x_range
    y_min_data, y_max_data = y_range
    for x_p, y_p in pixel_coords:
        data_x = x_min_data + (x_p / px_width) * (x_max_data - x_min_data)
        data_y = y_max_data - (y_p / px_height) * (y_max_data - y_min_data)
        data_coords.append((data_x, data_y))

    if num_points and num_points > 1:
        original_x = [p[0] for p in data_coords]
        original_y = [p[1] for p in data_coords]
        new_x = np.linspace(min(original_x), max(original_x), num_points)
        new_y = np.interp(new_x, original_x, original_y)
        return list(zip(new_x, new_y))

    return data_coords


def main_cli() -> None:
    """定义并处理命令行接口。"""
    parser = argparse.ArgumentParser(
        description='从绘制的线条图像中提取数据点。',
        formatter_class=argparse.RawTextHelpFormatter,
    )

    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('--file', dest='image_file', help='输入图像文件的路径。')
    input_group.add_argument(
        '--clipboard', action='store_true', help='使用系统剪贴板中的图像。'
    )

    parser.add_argument(
        '--xrange',
        nargs=2,
        type=float,
        required=True,
        metavar=('MIN', 'MAX'),
        help='X轴范围。',
    )
    parser.add_argument(
        '--yrange',
        nargs=2,
        type=float,
        required=True,
        metavar=('MIN', 'MAX'),
        help='Y轴范围。',
    )
    parser.add_argument(
        '--num-points', type=int, help='(可选) 将输出重采样到 N 个等距点。'
    )
    parser.add_argument('--output', help='(可选) 将输出保存为 CSV 文件的路径。')

    color_group = parser.add_mutually_exclusive_group()
    color_group.add_argument(
        '--target-hex', help="(可选) 线条的 HEX 颜色 (例如, '#FF8800')。"
    )
    color_group.add_argument(
        '--target-rgb',
        nargs=3,
        type=int,
        metavar=('R', 'G', 'B'),
        help='(可选) 线条的 RGB 颜色 (例如, 255 136 0)。',
    )

    parser.add_argument(
        '--color-tolerance',
        type=int,
        default=15,
        help='(可选) 颜色匹配的容忍度 (0-179)。默认值: 15。',
    )

    args = parser.parse_args()

    image_source: Optional[ImageSource] = None
    if args.clipboard:
        image_source = get_image_from_clipboard()
        if image_source is None:
            print('错误: 剪贴板中未找到有效图像。', file=sys.stderr)
            sys.exit(1)
    else:
        image_source = args.image_file

    if image_source is None:
        raise AssertionError('Expected an image source from CLI arguments')

    try:
        # 确定哪个颜色参数由用户提供，并直接传递
        color_arg = None
        if args.target_hex:
            color_arg = args.target_hex
        elif args.target_rgb:
            color_arg = tuple(args.target_rgb)

        data = extract_line_data(
            image_input=image_source,
            x_range=tuple(args.xrange),
            y_range=tuple(args.yrange),
            num_points=args.num_points,
            target_color=color_arg,
            color_tolerance=args.color_tolerance,
        )

        output_str = 'x,y\n' + ''.join([f'{x:.6f},{y:.6f}\n' for x, y in data])

        if args.output:
            with open(args.output, 'w') as f:
                f.write(output_str)
            print(f'✅ 成功处理了 {len(data)} 个点并保存到 {args.output}')
        else:
            print(output_str)

    except (ValueError, FileNotFoundError) as e:
        print(f'❌ 错误: {e}', file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main_cli()
