# Copyright (c) ModelScope Contributors. All rights reserved.
import base64
import math
import numpy as np
import os
import re
import requests
import torch
from io import BytesIO
from PIL import Image
from requests.adapters import HTTPAdapter
from typing import Any, Callable, List, TypeVar, Union
from urllib3.util.retry import Retry

from swift.utils import get_env_args

# >>> internvl
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def _build_transform(input_size):
    import torchvision.transforms as T
    from torchvision.transforms.functional import InterpolationMode
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform


def _find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio


def _dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    # calculate the existing image aspect ratio
    target_ratios = set((i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1)
                        if min_num <= i * j <= max_num)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = _find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    # calculate the target width and height
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    # resize the image
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = ((i % (target_width // image_size)) * image_size, (i // (target_width // image_size)) * image_size,
               ((i % (target_width // image_size)) + 1) * image_size, ((i //
                                                                        (target_width // image_size)) + 1) * image_size)
        # split the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images


# <<< internvl


def rescale_image(img: Image.Image, max_pixels: int) -> Image.Image:
    import torchvision.transforms as T
    width = img.width
    height = img.height
    if max_pixels is None or max_pixels <= 0 or width * height <= max_pixels:
        return img

    ratio = width / height
    height_scaled = math.sqrt(max_pixels / ratio)
    width_scaled = height_scaled * ratio
    return T.Resize((int(height_scaled), int(width_scaled)))(img)


_T = TypeVar('_T')


def _check_path(path: str) -> Union[str, None]:
    """If it is a path, return the string; if it is base64, return None."""
    MAX_PATH_HEURISTIC = 2000
    if len(path) > MAX_PATH_HEURISTIC:
        return
    if os.path.exists(path):
        return os.path.abspath(path)
    data = path
    ROOT_IMAGE_DIR = get_env_args('ROOT_IMAGE_DIR', str, None)
    if ROOT_IMAGE_DIR is not None:
        path = os.path.join(ROOT_IMAGE_DIR, path)
    path = os.path.abspath(os.path.expanduser(path))
    if os.path.exists(path):
        return path
    if data.startswith('data:'):
        return
    try:
        base64.b64decode(data)
        return
    except Exception:
        pass
    return data


def load_file(path: Union[str, bytes, _T]) -> Union[BytesIO, _T]:
    res = path
    if isinstance(path, str):
        path = path.strip()
        if path.startswith('http'):
            retries = Retry(total=3, backoff_factor=1, allowed_methods=['GET'])
            with requests.Session() as session:
                session.mount('http://', HTTPAdapter(max_retries=retries))
                session.mount('https://', HTTPAdapter(max_retries=retries))

                timeout = float(os.getenv('SWIFT_TIMEOUT', '20'))
                request_kwargs = {'timeout': timeout} if timeout > 0 else {}

                response = session.get(path, **request_kwargs)
                response.raise_for_status()
                content = response.content
                res = BytesIO(content)
        else:
            data = path
            path = _check_path(path)
            if path is None:
                # base64_str
                if data.startswith('data:'):
                    match_ = re.match(r'data:(.+?);base64,(.+)', data)
                    assert match_ is not None
                    data = match_.group(2)
                data = base64.b64decode(data)
                res = BytesIO(data)
            else:
                with open(path, 'rb') as f:
                    res = BytesIO(f.read())
    elif isinstance(path, bytes):
        res = BytesIO(path)
    return res


def load_image(image: Union[str, bytes, Image.Image]) -> Image.Image:
    image = load_file(image)
    if isinstance(image, BytesIO):
        image = Image.open(image)
    if image.mode != 'RGB':
        image = image.convert('RGB')
    return image


def load_batch(path_list: List[Union[str, None, Any, BytesIO]],
               load_func: Callable[[Any], _T] = load_image) -> List[_T]:
    res = []
    assert isinstance(path_list, (list, tuple)), f'path_list: {path_list}'
    for path in path_list:
        if path is None:  # ignore None
            continue
        res.append(load_func(path))
    return res


def load_video_hf(videos: List[str]):
    from transformers.video_utils import load_video
    res = []
    video_metadata = []
    for video in videos:
        if isinstance(video, (list, tuple)) and isinstance(video[0], str):
            # Case a: Video is provided as a list of image file names
            video = [np.array(load_image(image_fname)) for image_fname in video]
            video = np.stack(video)
            metadata = None
        else:
            # Case b: Video is provided as a single file path or URL or decoded frames in a np.ndarray or torch.tensor
            video_load_backend = get_env_args('video_load_backend', str, 'pyav')
            video, metadata = load_video(
                video,
                backend=video_load_backend,
            )
        res.append(video)
        video_metadata.append(metadata)
    return res, video_metadata


def _get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
    if bound:
        start, end = bound[0], bound[1]
    else:
        start, end = -100000, 100000
    start_idx = max(first_idx, round(start * fps))
    end_idx = min(round(end * fps), max_frame)
    seg_size = float(end_idx - start_idx) / num_segments
    frame_indices = np.array(
        [int(start_idx + (seg_size / 2) + np.round(seg_size * idx)) for idx in range(num_segments)])
    return frame_indices


def transform_image(image, input_size=448, max_num=12):
    transform = _build_transform(input_size=input_size)
    images = _dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
    pixel_values = [transform(image) for image in images]
    pixel_values = torch.stack(pixel_values)
    return pixel_values


def load_video_internvl(video: Union[str, bytes], bound=None, num_segments=32):
    from decord import VideoReader, cpu
    video_io = load_file(video)
    vr = VideoReader(video_io, ctx=cpu(0), num_threads=1)
    max_frame = len(vr) - 1
    fps = float(vr.get_avg_fps())

    images = []
    frame_indices = _get_index(bound, fps, max_frame, first_idx=0, num_segments=num_segments)
    for frame_index in frame_indices:
        images.append(Image.fromarray(vr[frame_index].asnumpy()).convert('RGB'))
    return images


def load_video_cogvlm2(video: Union[str, bytes]) -> np.ndarray:
    from decord import VideoReader, bridge, cpu
    video_io = load_file(video)
    bridge.set_bridge('torch')
    clip_end_sec = 60
    clip_start_sec = 0
    num_frames = get_env_args('num_frames', int, 24)
    decord_vr = VideoReader(video_io, ctx=cpu(0))
    duration = len(decord_vr)  # duration in terms of frames
    start_frame = int(clip_start_sec * decord_vr.get_avg_fps())
    end_frame = min(duration, int(clip_end_sec * decord_vr.get_avg_fps())) if \
        clip_end_sec is not None else duration
    frame_id_list = np.linspace(start_frame, end_frame - 1, num_frames, dtype=int)
    video_data = decord_vr.get_batch(frame_id_list)
    video_data = video_data.permute(3, 0, 1, 2)
    return video_data


def load_video_llava(video: Union[str, bytes]) -> np.ndarray:
    import av
    video_io = load_file(video)
    container = av.open(video_io)
    total_frames = container.streams.video[0].frames
    num_frames = get_env_args('num_frames', int, 16)
    indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    frames = []
    container.seek(0)
    start_index = indices[0]
    end_index = indices[-1]
    for i, frame in enumerate(container.decode(video=0)):
        if i > end_index:
            break
        if i >= start_index and i in indices:
            frames.append(frame)
    return np.stack([x.to_ndarray(format='rgb24') for x in frames])


def load_video_minicpmv_mplug_owl3(video: Union[str, bytes], max_num_frames):

    from decord import VideoReader, cpu  # pip install decord

    def uniform_sample(_l, _n):
        gap = len(_l) / _n
        idxs = [int(i * gap + gap / 2) for i in range(_n)]
        return [_l[i] for i in idxs]

    video_io = load_file(video)
    vr = VideoReader(video_io, ctx=cpu(0))
    sample_fps = round(vr.get_avg_fps() / 1)  # FPS
    frame_idx = [i for i in range(0, len(vr), sample_fps)]

    if len(frame_idx) > max_num_frames:
        frame_idx = uniform_sample(frame_idx, max_num_frames)
    frames = vr.get_batch(frame_idx).asnumpy()
    frames = [Image.fromarray(v.astype('uint8')) for v in frames]
    return frames


def _load_audio_librosa(audio: Union[str, bytes], sampling_rate: int, mono: bool = True):
    import librosa
    try:
        audio_io = load_file(audio)
        return librosa.load(audio_io, sr=sampling_rate, mono=mono)
    except Exception:
        if isinstance(audio, str) and audio.startswith(('http://', 'https://')):
            import audioread
            audio_io = audioread.ffdec.FFmpegAudioFile(audio)
        else:
            audio_io = _check_path(audio) if isinstance(audio, str) else audio
        return librosa.load(audio_io, sr=sampling_rate, mono=mono)


# ref: https://github.com/vllm-project/vllm/blob/v0.23.0/vllm/multimodal/audio.py#L169-L224
def _resample_audio_pyav(audio: np.ndarray, *, orig_sr: float, target_sr: float) -> np.ndarray:
    import av
    orig_sr_int = int(round(orig_sr))
    target_sr_int = int(round(target_sr))
    if orig_sr_int == target_sr_int:
        return audio
    if audio.ndim == 2:
        return np.stack([_resample_audio_pyav(ch, orig_sr=orig_sr, target_sr=target_sr) for ch in audio], axis=0)
    expected_len = int(math.ceil(audio.shape[-1] * target_sr_int / orig_sr_int))
    min_samples = 1024
    audio_f32 = np.asarray(audio, dtype=np.float32)
    if len(audio_f32) < min_samples:
        audio_f32 = np.pad(audio_f32, (0, min_samples - len(audio_f32)))
    audio_f32 = audio_f32.reshape(1, -1)
    resampler = av.AudioResampler(format='fltp', layout='mono', rate=target_sr_int)
    frame = av.AudioFrame.from_ndarray(audio_f32, format='fltp', layout='mono')
    frame.sample_rate = orig_sr_int
    out_frames = resampler.resample(frame)
    out_frames.extend(resampler.resample(None))
    result = np.concatenate([f.to_ndarray() for f in out_frames], axis=1).squeeze(0)
    return result[:expected_len]


# ref: https://github.com/vllm-project/vllm/blob/v0.23.0/vllm/multimodal/media/audio.py#L45-L160
def _load_audio_soundfile_pyav(path: Union[str, bytes, BytesIO], *, sr: float, mono: bool = True):
    """soundfile first, pyav fallback — same strategy as vLLM multimodal audio loader."""
    bad_sf_codes = {0, 1, 3, 4}
    if not isinstance(path, BytesIO):
        path = load_file(path)

    def _load_soundfile():
        import soundfile
        with soundfile.SoundFile(path) as f:
            native_sr = f.samplerate
            y = f.read(dtype='float32', always_2d=False).T

        if mono and y.ndim > 1:
            y = np.mean(y, axis=tuple(range(y.ndim - 1)))

        if sr is not None and sr != native_sr:
            y = _resample_audio_pyav(y, orig_sr=native_sr, target_sr=sr)
            return y, int(sr)
        return y, native_sr

    def _load_pyav():
        import av
        path.seek(0)
        with av.open(path) as container:
            if not container.streams.audio:
                raise ValueError('No audio stream found.')
            stream = container.streams.audio[0]
            stream.thread_type = 'AUTO'
            native_sr = stream.rate
            target_sr = sr or native_sr
            chunks = []
            needs_resampling = not math.isclose(float(target_sr), float(native_sr), rel_tol=0.0, abs_tol=1e-6)
            resampler = av.AudioResampler(format='fltp', layout='mono', rate=target_sr) if needs_resampling else None
            for frame in container.decode(stream):
                if needs_resampling:
                    for out_frame in resampler.resample(frame):
                        chunks.append(out_frame.to_ndarray())
                else:
                    chunks.append(frame.to_ndarray())
        if not chunks:
            raise ValueError('No audio found.')
        y = np.concatenate(chunks, axis=-1).astype(np.float32)
        if mono and y.ndim > 1:
            y = np.mean(y, axis=0)
        return y, target_sr

    try:
        return _load_soundfile()
    except ImportError:
        path.seek(0)
        return _load_pyav()
    except Exception as exc:
        import soundfile
        if not isinstance(exc, soundfile.LibsndfileError) or exc.code not in bad_sf_codes:
            raise
    path.seek(0)
    return _load_pyav()


def load_audio(
    audio: Union[str, bytes, BytesIO],
    sampling_rate: int,
    return_sr: bool = False,
    mono: bool = True,
):
    backend = get_env_args('swift_audio_load_backend', str, 'librosa')
    if backend == 'librosa':
        res = _load_audio_librosa(audio, sampling_rate, mono=mono)
    elif backend == 'soundfile_pyav':
        res = _load_audio_soundfile_pyav(audio, sr=sampling_rate, mono=mono)
    else:
        raise ValueError(f'Unknown audio load backend {backend!r}. Supported: librosa, soundfile_pyav')
    return res if return_sr else res[0]


def _resolve_video_local_path(path: Union[str, bytes]) -> tuple:
    """Return (local_path, is_temp_file). HTTP URLs and raw bytes are written to a temp file."""
    if isinstance(path, bytes) or (isinstance(path, str) and path.startswith('http')):
        import tempfile
        with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as f:
            f.write(load_file(path).read())
            return f.name, True
    checked = _check_path(path) if isinstance(path, str) else None
    return checked or path, False


def _video_to_ndarrays_local(local_path: str, num_frames: int = -1) -> np.ndarray:
    import cv2
    cap = cv2.VideoCapture(local_path)
    if not cap.isOpened():
        raise ValueError(f'Could not open video file {local_path}')
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames <= 0:
        cap.release()
        raise ValueError(f'Video file {local_path} has invalid or zero frame count: {total_frames}')
    if num_frames <= 0 or num_frames > total_frames:
        num_frames = total_frames
    frame_indices = set(np.linspace(0, total_frames - 1, num_frames, dtype=int))
    frames = []
    for idx in range(total_frames):
        if not cap.grab():
            break
        if idx in frame_indices:
            ret, frame = cap.retrieve()
            if ret:
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    if len(frames) < num_frames:
        raise ValueError(f'Could not read enough frames from video file {local_path} '
                         f'(expected {num_frames} frames, got {len(frames)})')
    return np.stack(frames)


def _video_get_metadata_local(local_path: str, num_frames: int = -1) -> dict:
    import cv2
    cap = cv2.VideoCapture(local_path)
    if not cap.isOpened():
        raise ValueError(f'Could not open video file {local_path}')
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames <= 0:
        cap.release()
        raise ValueError(f'Video file {local_path} has invalid or zero frame count: {total_frames}')
    fps = cap.get(cv2.CAP_PROP_FPS)
    duration = total_frames / fps if fps > 0 else 0
    cap.release()
    if num_frames <= 0 or num_frames > total_frames:
        num_frames = total_frames
    return {
        'total_num_frames': num_frames,
        'fps': duration / num_frames if num_frames else fps,
        'duration': duration,
        'video_backend': 'opencv',
        'frames_indices': list(range(num_frames)),
        'do_sample_frames': num_frames == total_frames,
    }


def load_vllm_video(path: Union[str, bytes], num_frames: int = -1) -> tuple:
    """Decode video frames + metadata for vLLM rollout; one download, temp file cleaned up."""
    local_path, is_temp = _resolve_video_local_path(path)
    try:
        return _video_to_ndarrays_local(local_path, num_frames), _video_get_metadata_local(local_path, num_frames)
    finally:
        if is_temp:
            try:
                os.remove(local_path)
            except OSError:
                pass


def load_video_valley(video: Union[str, bytes]):
    import decord
    from torchvision import transforms
    video_io = load_file(video)
    video_reader = decord.VideoReader(video_io)
    decord.bridge.set_bridge('torch')
    video = video_reader.get_batch(np.linspace(0, len(video_reader) - 1, 8).astype(np.int_)).byte()
    images = [transforms.ToPILImage()(image.permute(2, 0, 1)).convert('RGB') for image in video]
    return images


def load_video_ovis2(video_path, num_frames):
    from moviepy.editor import VideoFileClip
    with VideoFileClip(video_path) as clip:
        total_frames = int(clip.fps * clip.duration)
        if total_frames <= num_frames:
            sampled_indices = range(total_frames)
        else:
            stride = total_frames / num_frames
            sampled_indices = [
                min(total_frames - 1, int((stride * i + stride * (i + 1)) / 2)) for i in range(num_frames)
            ]
        frames = [clip.get_frame(index / clip.fps) for index in sampled_indices]
        frames = [Image.fromarray(frame, mode='RGB') for frame in frames]
    return frames


def load_video_ovis2_5(video_path, num_frames):
    from moviepy.editor import VideoFileClip
    with VideoFileClip(video_path) as clip:
        total_frames = int(clip.fps * clip.duration)
        indices = [int(i * total_frames / num_frames) for i in range(num_frames)]
        frames = [Image.fromarray(clip.get_frame(t)) for t in (idx / clip.fps for idx in indices)]
    return frames