import json
import os
import pandas as pd
import getpass

from PIL import Image
import io
import base64
from datetime import datetime

import use_vlm_ft_OpenAI


def get_img_base64_str(image_path):
    img = Image.open(image_path)
    buffered = io.BytesIO()
    img.save(buffered, format=img.format)
    img_base64_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return img_base64_str

if __name__ == "__main__":
    api_key = getpass.getpass("Enter your OpenAI API key: ")

    in_path_frame_images        = './.../rpp-765k_512'
    in_file_data_entering_test  = './.../test.parquet'
    
    in_model_name               = "OpenAI_FT_gpt-4o-2024-08-06"
    in_name_results_output      = 'OpenAI_FT_gpt-4o-2024-08-06'
    out_path_results            = './.../output'

    path_outputs = os.path.join(out_path_results, in_model_name)
    os.makedirs(path_outputs, exist_ok=True)

    os.environ["USED_MODEL"] = "gpt-4o-2024-08-06"

    dict_log = {}
    dict_log['model']           = in_model_name
    
    ft_model = "id-of-fine-tuned-model"
    #######################################################################
    #######################################################################
    df_result = pd.DataFrame(
        columns=['label', 'filename', \
        'brand', 'product_category', 'price', 'regular_price', 'relative_discount', 'absolute_discount', 'GTINs', 'weight_number', 'weight_unit', 'different_types'])
    df_result_cost = pd.DataFrame(
        columns=['label', 'filename']
    )

    df_test     = pd.read_parquet(in_file_data_entering_test, engine='pyarrow')
    df_test.reset_index(drop=True, inplace=True)

    start_index = 0
    output_file = f'{datetime.now().strftime("%Y%m%d_%H%M%S")}_{in_name_results_output}'
    for index, row in df_test.iloc[start_index:].iterrows():
        label       = str(row.label)
        filename    = row.filename

        dict_result = {}
        dict_result['label']    = label
        dict_result['filename'] = filename
        dict_result_cost = {}
        dict_result_cost['label']    = label
        dict_result_cost['filename'] = filename
            
        ############################################
        # PROMPT
        ############################################
        # IMAGE
        # image_path = os.path.join( in_path_frame_images, 'test', label, filename )
        image_path = os.path.join( in_path_frame_images, 'train', label, filename )
        query_image_base64 = get_img_base64_str(image_path)

        # TASK
        task = "Extract all targets."
        dict_log['prompt_task'] = task

        dict_log, dict_result, dict_result_cost = use_vlm_ft_OpenAI.do_request(
            api_key,
            ft_model,
            query_image_base64, 
            task, 
            dict_log,
            dict_result, 
            dict_result_cost,
        )
        
        df_result = pd.concat( [ df_result, pd.DataFrame.from_dict([dict_result]) ], ignore_index=True)
        df_result_cost = pd.concat( [ df_result_cost, pd.DataFrame.from_dict([dict_result_cost]) ], ignore_index=True)

        if index%100 == 0:
            df_result.to_parquet( os.path.join(path_outputs, output_file + '_' + str(index) + '.parquet'), index=False, engine='pyarrow')
            df_result_cost.to_parquet( os.path.join(path_outputs, output_file + '_costs_' +  str(index) + '.parquet'), index=False, engine='pyarrow')
            df_result = pd.DataFrame( columns=['label', 'filename', \
                        'brand', 'product_category', 'price', 'regular_price', 'relative_discount', 'absolute_discount', 'GTINs', 'weight_number', 'weight_unit', 'different_types'])
            df_result_cost = pd.DataFrame(columns=['label', 'filename'])
    
    df_result.to_parquet( os.path.join(path_outputs, output_file + '_' + str(index) + '.parquet'), index=False, engine='pyarrow')
    df_result_cost.to_parquet( os.path.join(path_outputs, output_file + '_costs_' +  str(index) + '.parquet'), index=False, engine='pyarrow')

    #######################################################################
    #######################################################################
    
    with open(os.path.join(path_outputs, in_name_results_output + '.json'), 'w') as json_file:
        json.dump(dict_log, json_file)