import time
import json
import ast
import io
from PIL import Image, ImageDraw, ImageFont, ImageColor
import xml.etree.ElementTree as ET
import numpy as np

import spaces
import torch
import gradio as gr
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor

### Load the model and helper functions ### 
model_path = "Qwen/Qwen2.5-VL-3B-Instruct"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    attn_implementation="sdpa"
).to(device)

processor = AutoProcessor.from_pretrained(model_path)
additional_colors = [colorname for (colorname, colorcode) in ImageColor.colormap.items()]

def parse_json(json_output):
    lines = json_output.splitlines()
    for i, line in enumerate(lines):
        if line.startswith("```json"):
            json_output = "\n".join(lines[i+1:])
            json_output = json_output.split("```")[0]
            break
    return json_output

def plot_bounding_boxes(im, bounding_boxes, input_width, input_height):
    img = im
    width, height = img.size
    draw = ImageDraw.Draw(img)
    colors = ['red', 'green', 'blue', 'yellow', 'orange', 'pink', 'purple', 'cyan', 'magenta', 'lime'] + additional_colors

    bounding_boxes = parse_json(bounding_boxes)
    try:
        font = ImageFont.load_default(size=16.0)
    except:
        font = ImageFont.load_default()

    try:
      json_output = ast.literal_eval(bounding_boxes)
    except Exception as e:
      end_idx = bounding_boxes.rfind('"}') + len('"}')
      truncated_text = bounding_boxes[:end_idx] + "]"
      try:
          json_output = ast.literal_eval(truncated_text)
      except:
          json_output = []

    for i, bounding_box in enumerate(json_output):
      color = colors[i % len(colors)]
      try:
          coords = bounding_box.get("bbox_2d", bounding_box.get("box_2d", [0,0,0,0]))
          abs_y1 = int(coords[1]/input_height * height)
          abs_x1 = int(coords[0]/input_width * width)
          abs_y2 = int(coords[3]/input_height * height)
          abs_x2 = int(coords[2]/input_width * width)
      except:
          continue

      if abs_x1 > abs_x2: abs_x1, abs_x2 = abs_x2, abs_x1
      if abs_y1 > abs_y2: abs_y1, abs_y2 = abs_y2, abs_y1

      draw.rectangle(((abs_x1, abs_y1), (abs_x2, abs_y2)), outline=color, width=4)

      # --- HIỂN THỊ CALO ---
      if "label" in bounding_box:
        display_text = bounding_box["label"]
        if "nutrition" in bounding_box:
            cals = bounding_box["nutrition"].get("calories_kcal", "?")
            display_text += f" ({cals} kcal)"
        
        try:
            bbox = draw.textbbox((abs_x1 + 8, abs_y1 + 6), display_text, font=font)
            draw.rectangle(bbox, fill="black")
        except AttributeError:
            pass 
        draw.text((abs_x1 + 8, abs_y1 + 6), display_text, fill=color, font=font)
    return img

def resize_image_to_max_dimension(image, max_dimension=1024):
    width, height = image.size
    if width <= max_dimension and height <= max_dimension: return image
    if width >= height:
        new_width = max_dimension
        new_height = int((height / width) * max_dimension)
    else:
        new_height = max_dimension
        new_width = int((width / height) * max_dimension)
    return image.resize((new_width, new_height), Image.LANCZOS)

def inference(img_url, prompt, system_prompt="You are a nutrition assistant. Analyze food images, locate items, and estimate nutrition.", max_new_tokens=1024):
  if isinstance(img_url, str):
    image = Image.open(img_url)
  else:
    image = resize_image_to_max_dimension(img_url)

  messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": [{"type": "text", "text": prompt}, {"image": img_url}]}
  ]
  text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
  inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt").to(device)
  output_ids = model.generate(**inputs, max_new_tokens=1024)
  generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
  output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
  
  input_height = inputs['image_grid_thw'][0][1]*14
  input_width = inputs['image_grid_thw'][0][2]*14
  return output_text[0], input_height, input_width

def numpy_to_pil(numpy_array):
    if numpy_array.dtype != np.uint8:
        if numpy_array.dtype in [np.float32, np.float64] and numpy_array.max() <= 1.0:
            numpy_array = (numpy_array * 255).astype(np.uint8)
        else:
            numpy_array = numpy_array.astype(np.uint8)
    return Image.fromarray(numpy_array)

### Prompt section (Đã được rút gọn) ###
prompt_thinking = """Outline the bounding box coordinates and names of each unique edible food and drink item and output all the coordinates in JSON format.
Crucially, for EACH item, you must also estimate its nutritional value per serving (calories, protein, carbs, fat) and include it in the JSON under the "nutrition" key.
Start with <think></think> tags. Output format: [{"bbox_2d": ..., "label": ..., "nutrition": {"calories_kcal": ..., "protein_g": ..., "carbs_g": ..., "fat_g": ...}}...].

<examples>
<example_1>
<think>
The image contains a pint of dark beer and a green salad. I will create bounding box coordinates and estimate macros.
</think>
```json
[
  {"bbox_2d": [27, 271, 365, 531], "label": "Pint of Guinness", "nutrition": {"calories_kcal": 210, "protein_g": 1.5, "carbs_g": 18, "fat_g": 0}},
  {"bbox_2d": [343, 300, 609, 711], "label": "Green Salad", "nutrition": {"calories_kcal": 120, "protein_g": 3, "carbs_g": 10, "fat_g": 8}}
]
</example_1>

"""

prompt_no_thinking = """Outline the bounding box coordinates and names of each unique edible food and drink item and output all the coordinates in JSON format.
Crucially, for EACH item, you must also estimate its nutritional value per serving (calories, protein, carbs, fat) and include it in the JSON under the "nutrition" key.
Output format: [{"bbox_2d": ..., "label": ..., "nutrition": {"calories_kcal": ..., "protein_g": ..., "carbs_g": ..., "fat_g": ...}}...].

@spaces.GPU(duration=120)
def infer_on_image(input_image):
image_1 = numpy_to_pil(numpy_array=input_image)
start_time_thinking = time.time()
response, input_height, input_width = inference(img_url=image_1, prompt=prompt_thinking)
total_time_thinking = round(time.time() - start_time_thinking, 4)
output_image_1 = plot_bounding_boxes(image_1, response, input_width, input_height)

image_2 = numpy_to_pil(numpy_array=input_image)
start_time_no_thinking = time.time()
response_no_thinking, input_height, input_width = inference(img_url=image_2, prompt=prompt_no_thinking)
total_time_no_thinking = round(time.time() - start_time_no_thinking, 4)
output_image_2 = plot_bounding_boxes(image_2, response_no_thinking, input_width, input_height)

return output_image_1, response, total_time_thinking, output_image_2, response_no_thinking, total_time_no_thinking

demo = gr.Interface(fn=infer_on_image,
inputs=gr.Image(label="Input image"),
outputs=[gr.Image(label="Image w/ thinking tags"),
gr.Text(label="Raw output w/ thinking tags"),
gr.Text(label="Inference time w/ thinking tags"),
gr.Image(label="Image w/o thinking tags"),
gr.Text(label="Raw output w/o thinking tags"),
gr.Text(label="Inference time w/o thinking tags")],
title="Qwen2.5-VL Food & Nutrition Detection 👁️🍔",
description="Nhận diện món ăn, vẽ Bounding Box và Ước tính Dinh dưỡng.",
cache_examples=False)

demo.launch(debug=True)