from PIL import Image, ImageDraw
import random
import requests
import uuid
import base64
import io
import json
from json_repair import repair_json
from tqdm import tqdm
import math
from typing import Sequence
import time
import traceback
import warnings

prompt_pairwise = """
You're labeling some preference data to train a subject-driven image generating model, where 
1. The first given image is the original one (reference image), and the main subject "{subject}" is marked by a red bounding box.
2. The second and third images are the generated results
3. The generation is performed under this text instruction: "{instruction}".

You should choose one of the two generated results based on these criterion:
1. Composition plausibility: The subject should be seamlessly and naturally integrated into the new context, with coherent lighting, perspective, and style.
2. Subject Consistency: the main subject should not deviate too much from that in the reference image, except indicated by the text instruction or changes in lighting, angle, size, position, etc.
3. Absence of Irrelevant Artifacts: The generated image should not retain any elements from the reference image that are unrelated to the subject (e.g., background remnants, unrelated objects, or visual clutter).
4. Overall image quality: The image should be free from visual artifacts, distortions, or unnatural features that detract from its realism and aesthetic appeal.

Return a single letter for the better image: "A" if the second image is better, or "B" if the third image is better.
"""

def img_to_base64(image: str | Image.Image) -> str:
    image = Image.open(image) if not isinstance(image, Image.Image) else image
    image = image.convert('RGB')
    img_bytes = io.BytesIO()
    image.save(img_bytes, format='JPEG')
    image_bytes = img_bytes.getvalue()
    encoded_string = base64.b64encode(image_bytes)
    img_base64_str = encoded_string.decode("utf-8")

    return img_base64_str


def resize_keep_h(im: Image.Image, target_h: int) -> Image.Image:
    w, h = im.size
    scale = target_h / h
    return im.resize((int(w * scale), target_h), Image.LANCZOS)


class VlmClient:
    def __init__(self, model, endpoint, token):
        self.endpoint = endpoint
        self.model = model
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
        }

    def query(self, images, prompt):
        if not isinstance(images, list):
            raise ValueError()
        messages = [
            {
                "role": "user",
                "content": [ 
                    {"type": "text", "text": prompt},
                ]
            }]
        for image in images:
            messages[0]['content'].append({
                "type": "image_url", 
                "image_url": {
                    "url": f"data:image/jpeg;base64,{img_to_base64(image)}",
                    "detail": "high" 
                }
            })
        json_data = {
            "model": self.model,
            "messages": messages,
            "temperature": 1.0,
            "top_p": 1.0,
            "max_tokens": 1000,
            "stream": False,
            "logprobs": True,
            "top_logprobs": 4
        }

        resp = requests.post(self.endpoint, headers=self.headers, json=json_data, stream=True)
        try:
            resp = resp.json()['choices'][0]
        except Exception as e:
            print("Error in response:", resp.text)
            raise e
        return resp


class VlmPairwiseScorer:
    def __init__(self):
        self.st = time.time()
        self.nfail = 0
        model = ""
        endpoint = ""
        token = ""
        self.client = VlmClient(model=model, endpoint=endpoint, token=token)

    @staticmethod
    def token_prune(token):
        token = token.strip()
        token = token.replace(".", "").replace(",", "").replace("'", "").replace('"', "")
        return token

    def __call__(self, ref: Image.Image, images: Sequence[Image.Image], instruction: str, subject: str) -> Sequence[float]:
        wins = [0 for _ in images]
        ref = resize_keep_h(ref.convert("RGB"), min(ref.height, 512))
        images = [resize_keep_h(im.convert("RGB"), min(im.height, 512)) for im in images]
        npair = len(images) * (len(images) - 1) // 2
        for i, img1 in enumerate(images):
            for j in range(i + 1, len(images)):
                try:
                    img2 = images[j]
                    swap = random.random() < 0.5
                    if swap:
                        imgA, imgB = img2, img1
                    else:
                        imgA, imgB = img1, img2

                    response = self.client.query([ref, imgA, imgB], prompt_pairwise.format(subject=subject, instruction=instruction))

                    choice = response['message']['content']
                    assert choice in ["A", "B"], f"Unexpected choice: {choice}"

                    conf = {"A": -1.0, "B": -1.0}
                    logprobs = response['logprobs']['content']
                    logp = None
                    for ilogp in logprobs:
                        if self.token_prune(ilogp["token"]) in ["A", "B"]:
                            logp = ilogp
                            break
                    if logp is not None:
                        for topp in logp['top_logprobs']:
                            token = self.token_prune(topp["token"])
                            if token in ["A", "B"]:
                                conf[token] = math.exp(topp["logprob"])
                    if conf["A"] < 0 or conf["B"] < 0:
                        print(f"({i}, {j}): No valid logprobs found, defaulting to binary confidence.")
                        conf[choice] = 1.0
                        conf["A" if choice == "B" else "B"] = 0.0
                    total = sum(conf.values())
                    for k in conf:
                        conf[k] /= total

                    choice2idx = {"A": i if not swap else j, "B": j if not swap else i}
                    for choice, idx in choice2idx.items():
                        wins[idx] += conf[choice]
                except Exception as e:
                    print(f"VLMPairwiseScorer: Error processing pair ({i}, {j}): {e}")
                    traceback.print_exc()
                    wins[i] += 0.5
                    wins[j] += 0.5
                    self.nfail += 1
        elapsed = time.time() - self.st
        fail_per_min = self.nfail / (elapsed / 60.0 + 1e-5)
        if fail_per_min > 10:
            warnings.warn(f"Too many failures in VLM queries: {self.nfail} failures over {elapsed:.1f} seconds ({fail_per_min:.1f} per minute)")
        if elapsed // 60 > 10:
            print(f"VLMPairwiseScorer: {self.nfail} failures ({fail_per_min:.1f} per minute) in {elapsed/60:.1f} minutes")
            self.st = time.time()
            self.nfail = 0

        assert abs(sum(wins) - npair) < 1e-5, f"Sum of wins {sum(wins)} != number of pairs {npair}"
        return wins


if __name__ == "__main__":
    scorer = VlmPairwiseScorer()
