import os
import time
import logging
import json
import traceback

import torch
import requests
from PIL import Image
from transformers import AutoModelForCausalLM, LlamaTokenizer, AutoTokenizer
from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch

DATASET_PATH = '../../evaluation/gts/det/semantics/union3_test.json'
IMG_DIR = '../../dataset/data/coco_det/images/semantics/'
OUTPUT_PATH = './output1.json'

MODEL_PATH = "THUDM/cogvlm-chat-hf"
# MODEL_PATH = "zai-org/cogvlm-grounding-generalist-hf"
TOKENIZER_PATH = "lmsys/vicuna-7b-v1.5"
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

tokenizer = LlamaTokenizer.from_pretrained(TOKENIZER_PATH)
torch_type = torch.bfloat16

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch_type,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(DEVICE).eval()


def process_image(image_path):
    # chat example
    query = '''--- Task ---
This is a screenshot of a VR game with a size of 960*540. Please identify all interactable objects on the screenshot, describe what they are, locate them with bounding box in the image and give how confident you are about the result, ranging from 0 to 1.
--- Bounding Box Format ---
[x, y, width, height], where x, y are the coordinates of the top-left corner of the box, and width, height are the width and height of the box. The unit of x, y, width, height are pixel.
x coordinate is the horizontal distance from the left edge of the image, and y coordinate is the vertical distance from the top edge of the image.
--- Output Format ---
Output in JSON format: [{"name": object_name1, "bbox": bounding_box1, "confidence": confidence1}, {"name": object_name2, "bbox": bounding_box2, "confidence": confidence2}].
DO NOT OUTPUT any other content besides JSON.'''
    image = Image.open(image_path).convert('RGB')
    query = "USER: {} ASSISTANT:".format(query)

    input_by_model = model.build_conversation_input_ids(tokenizer, query=query, history=[], images=[image])

    inputs = {
        'input_ids': input_by_model['input_ids'].unsqueeze(0).to(DEVICE),
        'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(DEVICE),
        'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(DEVICE),
        'images': [[input_by_model['images'][0].to(DEVICE).to(torch_type)]] if image is not None else None,
    }
    if 'cross_images' in input_by_model and input_by_model['cross_images']:
        inputs['cross_images'] = [[input_by_model['cross_images'][0].to(DEVICE).to(torch_type)]]

    # add any transformers params here.
    gen_kwargs = {"max_length": 2048,
                    "do_sample": False} # "temperature": 0.9
    with torch.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        response = tokenizer.decode(outputs[0])
        response = response.split("</s>")[0]
    response = response.strip().strip('```json').strip('```').strip()
    return response



def main():
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
    start_time = time.time()
    logging.info('Started at %s', time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))

    with open(DATASET_PATH, 'r') as f:
        dataset = json.load(f)

    if os.path.exists(OUTPUT_PATH):
        with open(OUTPUT_PATH, 'r') as f:
            output = json.load(f)
    else:
        output = {}

    # progress = tqdm(dataset['images'], total=len(dataset['images']), leave=True, position=0)
    # for img in progress:
    for img in dataset['images']:
        img_id = img['id']
        img_name = img['file_name']
        if str(img_id) in output.keys() and not output[str(img_id)] == '':
            logging.info(f"Skipping {img_name}")
            continue
        logging.info(f"Processing {img_name}")
        img_path = os.path.join(IMG_DIR, img_name)
        # progress.set_description(f"Processing {img_name}")
        objects = ''
        try:
            objects = process_image(img_path)
        except KeyboardInterrupt:
            logging.info("Process interrupted by user. Exiting.")
            break
        except Exception as e:
            logging.error(f"Error processing image {img_name}: {e}")
            objects = ''
        finally:
            output[str(img_id)] = objects
            with open(OUTPUT_PATH, 'w') as f:
                json.dump(output, f, indent=4)

    logging.info('Completed at %s', time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))
    logging.info(f"Total time taken: {time.time() - start_time} seconds")


if __name__ == "__main__":
    main()