import b2sdk.v2 as b2 #Backblaze img2img upload bucket
import base64
import shutil
import os
import requests
import io
from PIL import Image
import numpy as npzipfile
import zipfile

import tempfile
import json 
import gradio as gr 

def update_model_dicts(traning_finnal,token_string,style_json="model_dict.json"):

  print(traning_finnal,token_string)
  current_style_dict=json.load(open(style_json,"r"))
  current_style_dict[token_string]=traning_finnal
  with open(style_json, "w") as json_file:
    json.dump(current_style_dict, json_file, indent=4)
  json_file.close()
  # Return the updated dictionary keys for updating the Dropdown
  return list(current_style_dict.keys())

def update_dropdown(traning_finnal, token_string):
    updated_keys = update_model_dicts(traning_finnal, token_string)
    return gr.Dropdown.update(choices=updated_keys)

def add_to_prompt(existing_prompt, new_prompt):
    if existing_prompt:
        return f"{existing_prompt}, {new_prompt}"
    else:
        return new_prompt


def numpy_to_base64(image_np):
  """Converts a numpy image to base64 string."""
  img = Image.fromarray(image_np)
  buffered = io.BytesIO()
  img.save(buffered, format="PNG")
  img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
  return "data:image/png;base64,"+img_str
    
def image_to_base64(img):
  buffered = io.BytesIO()
  img.save(buffered, format="PNG")
  img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
  return "data:image/png;base64,"+img_str


def create_zip(files,captions,trigger):
    #Caption processing
    captions=captions.split("\n")
    #cute files and "tags:"
    captions= [cap.split("file:")[0][5:] for cap in captions]
    print("files",len(files),"captions",len(captions))
    #assert len(files)==len(captions) , "File amount does not equal the captions amount please check"
    temp_dir="./datasets/"
    os.makedirs(temp_dir,exist_ok=True)
    
    zip_path = os.path.join(temp_dir, f"training_data_{trigger}.zip")
    if os.path.exists(zip_path):
       os.remove(zip_path)

    with zipfile.ZipFile(zip_path, "w") as zip_file:
      for i, file in enumerate(files):
          # Add image to zip
          image_name = f"image_{i}.jpg"
          print(file)
          zip_file.write(file, image_name)
          # Add caption to zip
          caption_name = f"image_{i}.txt"
          caption_content = captions[i] +f", {trigger}"
          zip_file.writestr(caption_name, caption_content)
    return zip_path

def BB_uploadfile(local_file,file_name,BB_bucket_name,FRIENDLY_URL=True):
    info = b2.InMemoryAccountInfo()
    b2_api = b2.B2Api(info)
    #print(application_key_id,application_key)
    application_key_id = os.getenv("BB_KeyID")
    application_key = os.getenv("BB_AppKey")
    b2_api.authorize_account("production", application_key_id, application_key)
    BB_bucket=b2_api.get_bucket_by_name(BB_bucket_name)
    BB_defurl="https://f005.backblazeb2.com/file/"

    metadata = {"key": "value"}
    uploaded_file = BB_bucket.upload_local_file(
            local_file=local_file,
            file_name=file_name,
            file_infos=metadata,
    )
    img_url=b2_api.get_download_url_for_fileid(uploaded_file.id_)
    if FRIENDLY_URL: #Get friendly URP
        img_url=BB_defurl+BB_bucket_name+"/"+file_name
    print("backblaze", img_url)
    return img_url
#file="/content/training_data.zip"