# ADOBE CONFIDENTIAL
# Copyright 2025 Adobe
# All Rights Reserved.
# NOTICE: All information contained herein is, and remains
# the property of Adobe and its suppliers, if any. The intellectual
# and technical concepts contained herein are proprietary to Adobe
# and its suppliers and are protected by all applicable intellectual
# property laws, including trade secret and copyright laws.
# Dissemination of this information or reproduction of this material
# is strictly forbidden unless prior written permission is obtained
# from Adobe.

from __future__ import annotations
import torch
from .token_constants import DEFAULT_IMAGE_PAD_TOKEN, DEFAULT_IMAGE_TOKEN
import numpy as np
from PIL import ImageOps

def pad_to_square_and_resize(img, size=1024):
    padded_img = ImageOps.pad(img, (max(img.size), max(img.size)), color=(0, 0, 0))
    resized_img = padded_img.resize((size, size))
    return resized_img

def process_images(images, image_processor, model_cfg):
    result = []
    for image in images:
        pixels = image_processor(images=[image], return_tensors='pt').pixel_values
        (h, w) = pixels.shape[-2:]
        result.append((pixels, torch.tensor([[1, h // 16, w // 16]], dtype=torch.long)))
    return result

def tokenizer_image_token(prompt, tokenizer, image_token_index=-200, return_tensors=None, constant_version=None):
    prompt = prompt.replace('<image_padding>', '<SPECIAL_23>')
    ids = torch.tensor(tokenizer(prompt).input_ids, dtype=torch.long)
    ids[ids == 23] = -500
    return ids if return_tensors == 'pt' else ids.tolist()

def compute_num_image_tokens(x, image_aspect_ratio, image_processor):
    grid = x[0][1]
    return (int(grid[0, 1]) // 2, int(grid[0, 2]) // 2)

def fix_one_prompt(prompt, config, n_image_tokens):
    image_idx = 0
    while DEFAULT_IMAGE_TOKEN in prompt:
        newline_token = config.mm_native_newline_token
        if newline_token:
            if len(n_image_tokens[image_idx]) == 2:
                (hh, ww) = n_image_tokens[image_idx]
                row = DEFAULT_IMAGE_PAD_TOKEN * ww + newline_token
                rows = row * hh
                rows = rows.rstrip(newline_token)
                prompt = prompt.replace('<image>', rows, 1)
            elif len(n_image_tokens[image_idx]) == 3:
                (tt, hh, ww) = (int(value) for value in n_image_tokens[image_idx])
                row = DEFAULT_IMAGE_PAD_TOKEN * ww + newline_token
                per_image = row * hh + newline_token + newline_token
                images = (per_image * tt).rstrip(newline_token + newline_token)
                prompt = prompt.replace('<image>', images, 1)
            else:
                raise NotImplementedError()
        else:
            image_tokens = DEFAULT_IMAGE_PAD_TOKEN * np.prod(n_image_tokens[image_idx])
            prompt = prompt.replace('<image>', image_tokens, 1)
        image_idx += 1
    return prompt