"""1. 图像预处理节点：获取图片信息，直接使用原始URL"""
import logging
import urllib.request
import time
from io import BytesIO
from typing import Tuple, Optional
from langchain_core.runnables import RunnableConfig
from langgraph.runtime import Runtime
from coze_coding_utils.runtime_ctx.context import Context
from PIL import Image

from graphs.state import (
    ImagePreprocessInput,
    ImagePreprocessOutput,
    ImageInfo
)

logger = logging.getLogger(__name__)

# 默认图片尺寸（用于降级处理）
DEFAULT_IMAGE_SIZE = (1000, 1400)

# 超时配置（秒）
IMAGE_DOWNLOAD_TIMEOUT = 30  # 单次下载超时
MAX_RETRIES = 2  # 最大重试次数（减少重试）


def get_image_info_with_retry(image_url: str, max_retries: int = MAX_RETRIES, timeout: int = IMAGE_DOWNLOAD_TIMEOUT) -> Tuple[int, int, int]:
    """
    获取图片尺寸信息（带重试机制，有最大总时间限制）
    
    Args:
        image_url: 图片URL
        max_retries: 最大重试次数（默认2次）
        timeout: 单次请求超时时间（默认30秒）
        
    Returns:
        (width, height, dpi)
    """
    last_error = None
    total_start_time = time.time()
    MAX_TOTAL_TIME = 60  # 总时间不超过60秒
    
    for attempt in range(max_retries):
        # 检查总时间
        if time.time() - total_start_time > MAX_TOTAL_TIME:
            logger.warning(f"Total time exceeded {MAX_TOTAL_TIME}s, giving up")
            break
            
        try:
            # 下载图片（带超时）
            with urllib.request.urlopen(image_url, timeout=timeout) as response:
                img_data = response.read()
            
            # 检查数据大小
            if len(img_data) < 100:
                raise ValueError(f"Image data too small: {len(img_data)} bytes")
            
            img = Image.open(BytesIO(img_data))
            
            # 转换为RGB模式（处理PNG透明通道）
            if img.mode in ('RGBA', 'P'):
                img = img.convert('RGB')
            
            width, height = img.size
            
            # 验证尺寸有效性
            if width <= 0 or height <= 0:
                raise ValueError(f"Invalid image size: {width}x{height}")
            
            dpi = img.info.get('dpi', (72, 72))
            if isinstance(dpi, tuple):
                dpi = dpi[0] if dpi[0] > 0 else 72
            else:
                dpi = 72
            
            elapsed = time.time() - total_start_time
            logger.info(f"Got image info: {width}x{height} in {elapsed:.1f}s (attempt {attempt + 1})")
            return width, height, dpi
            
        except Exception as e:
            last_error = e
            elapsed = time.time() - total_start_time
            logger.warning(f"Attempt {attempt + 1}/{max_retries} failed in {elapsed:.1f}s: {str(e)[:100]}")
            
            # 只有非最后一次才等待
            if attempt < max_retries - 1:
                wait_time = 1  # 固定等待1秒
                time.sleep(wait_time)
    
    # 所有重试都失败，抛出错误
    raise RuntimeError(f"图片获取失败（{max_retries}次重试）: {str(last_error)[:100]}")


def image_preprocess_node(
    state: ImagePreprocessInput,
    config: RunnableConfig,
    runtime: Runtime[Context]
) -> ImagePreprocessOutput:
    """
    title: 图像预处理
    desc: 获取图片尺寸信息，直接使用原始图片URL用于AI识别，不旋转不缩放不上传。支持重试和降级处理。
    integrations: 
    """
    ctx = runtime.context
    
    # 获取原始图片URL
    original_image_url = state.homework_image.url
    
    try:
        # 获取图片尺寸信息（带重试）
        width, height, dpi = get_image_info_with_retry(original_image_url)
        logger.info(f"原始图片尺寸：宽{width} x 高{height}")
        
        return ImagePreprocessOutput(
            image_info=ImageInfo(
                width=width,
                height=height,
                dpi=dpi
            ),
            image_url=original_image_url
        )
        
    except Exception as e:
        # 降级处理：使用默认尺寸，但仍然继续处理
        logger.error(f"Failed to get image info after retries, using fallback: {e}")
        
        # 返回默认尺寸，让后续节点能够继续处理
        return ImagePreprocessOutput(
            image_info=ImageInfo(
                width=DEFAULT_IMAGE_SIZE[0],
                height=DEFAULT_IMAGE_SIZE[1],
                dpi=72
            ),
            image_url=original_image_url
        )
