import os
from aiohttp import web
import server
import hashlib
import folder_paths
import uuid
import time
import json
from .libs.prompt_parser import get_learned_conditioning_prompt_schedules
from .sdfx_path import get_gallery_path, get_workflows_path, get_templates_path

gallery_path = get_gallery_path()
workflows_path = get_workflows_path()
templates_path = get_templates_path()

def get_media_url_from_params(url_path, filename, type, gallery):
  return url_path + '/sdfx/media?filename=' + filename + '&type=' + type + '&gallery=' + gallery

def get_dir_by_model_id_and_type(model_id, dir_type):
  # === Todo handle other media types ===
  if dir_type is None or dir_type not in ["image"]:
      dir_type = "image"

  if dir_type == "image":
      type_dir = os.path.join(gallery_path, model_id, "image")

  return type_dir, dir_type

def media_upload(post, url_path):
  image = post.get("file")
  overwrite = post.get("overwrite")
  model_id = post.get("modelId")

  image_upload_type = post.get("type")
  upload_dir, image_upload_type = get_dir_by_model_id_and_type(model_id, image_upload_type)

  if image and image.file:
    filename = image.filename
    if not filename:
      return web.Response(status=400)

    filepath = os.path.abspath(os.path.join(upload_dir, filename))

    if os.path.commonpath((upload_dir, filepath)) != upload_dir:
      return web.Response(status=400)

    if not os.path.exists(upload_dir):
      os.makedirs(upload_dir)

    split = os.path.splitext(filename)

    if overwrite is not None and (overwrite == "true" or overwrite == "1"):
      pass
    else:
      i = 1
      while os.path.exists(filepath):
        filename = f"{split[0]} ({i}){split[1]}"
        filepath = os.path.join(upload_dir, filename)
        i += 1

    with open(filepath, "wb") as f:
      f.write(image.file.read())
    return web.json_response({"name" : filename, "type": image_upload_type, "path": get_media_url_from_params(url_path, filename, image_upload_type, model_id)})
  else:
    return web.Response(status=400)

def get_domain_from_request(request):
  return request.scheme + '://' + request.host

def get_medias_from_model_id(media_id, url_path):
  media_type = 'image'
  # === Todo handle other media types ===
  media_folder = os.path.join(gallery_path, media_id, media_type)
  if os.path.exists(media_folder) and os.path.isdir(media_folder):
    return [{"name": image, "type": media_type, "path": get_media_url_from_params(url_path, image, media_type, media_id)} for image in os.listdir(media_folder)]

  return []

def add_sdfx_routes():
  # === Comfy determines the path of a checkpoint by its name ===
  # === If two different paths contain two checkpoints with the same name, there is an arbitration done to choose the appropriate path ===
  @server.PromptServer.instance.routes.get("/sdfx/model/list")
  def get_models(request):
    url_path = get_domain_from_request(request)
    models = []
    cns = folder_paths.get_filename_list('checkpoints')
    print(cns)
    for cn in cns:
        fp = folder_paths.get_full_path("checkpoints", cn)
        media_id = hashlib.md5(fp.encode()).hexdigest()
        models.append({
            "modelId": media_id,
            "name": cn,
            "type": "checkpoint",
            "gallery": get_medias_from_model_id(media_id, url_path)
        })

    return web.json_response(models)

  # === Single API to upload all kind of assets ===
  @server.PromptServer.instance.routes.post("/sdfx/media")
  async def upload_media(request):
    post = await request.post()
    url_path = get_domain_from_request(request)
    return media_upload(post, url_path)
  

  @server.PromptServer.instance.routes.get("/sdfx/media")
  async def view_media(request):
    if "filename" not in request.rel_url.query:
      return web.Response(status=400)
    if "type" not in request.rel_url.query:
      return web.Response(status=400)
    if "gallery" not in request.rel_url.query:
      return web.Response(status=400)

    filename = request.rel_url.query["filename"]
    type = request.rel_url.query["type"]
    gallery = request.rel_url.query["gallery"]

    if filename[0] == '/' or '..' in filename:
      return web.Response(status=400)
    if type[0] == '/' or '..' in type:
      return web.Response(status=400)
    if gallery[0] == '/' or '..' in gallery:
      return web.Response(status=400)

    file = os.path.join(gallery_path, gallery, type, filename)

    if os.path.isfile(file):
      if 'preview' in request.rel_url.query:
        with Image.open(file) as img:
          preview_info = request.rel_url.query['preview'].split(';')
          image_format = preview_info[0]
          if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''):
            image_format = 'webp'

          quality = 90
          if preview_info[-1].isdigit():
            quality = int(preview_info[-1])

          buffer = BytesIO()
          if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb':
            img = img.convert("RGB")
          img.save(buffer, format=image_format, quality=quality)
          buffer.seek(0)

          return web.Response(body=buffer.read(), content_type=f'image/{image_format}',
                              headers={"Content-Disposition": f"filename=\"{filename}\""})

      if 'channel' not in request.rel_url.query:
        channel = 'rgba'
      else:
        channel = request.rel_url.query["channel"]

      if channel == 'rgb':
        with Image.open(file) as img:
          if img.mode == "RGBA":
            r, g, b, a = img.split()
            new_img = Image.merge('RGB', (r, g, b))
          else:
            new_img = img.convert("RGB")

          buffer = BytesIO()
          new_img.save(buffer, format='PNG')
          buffer.seek(0)

          return web.Response(body=buffer.read(), content_type='image/png', headers={"Content-Disposition": f"filename=\"{filename}\""})

      elif channel == 'a':
        with Image.open(file) as img:
          if img.mode == "RGBA":
            _, _, _, a = img.split()
          else:
            a = Image.new('L', img.size, 255)

          # alpha img
          alpha_img = Image.new('RGBA', img.size)
          alpha_img.putalpha(a)
          alpha_buffer = BytesIO()
          alpha_img.save(alpha_buffer, format='PNG')
          alpha_buffer.seek(0)

          return web.Response(body=alpha_buffer.read(), content_type='image/png', headers={"Content-Disposition": f"filename=\"{filename}\""})
      else:
        return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""})

    return web.Response(status=400)

  @server.PromptServer.instance.routes.get("/sdfx/prompt_parser")
  async def prompt_parser(request):
    if "prompt" in request.rel_url.query:
      prompt = request.rel_url.query["prompt"]
      env = request.rel_url.query['env'] if 'env' in request.rel_url.query else 'prod'
      g = lambda p: get_learned_conditioning_prompt_schedules([p], env)[0]
      return web.json_response({'prompt' : prompt, 'data' : g(prompt)})
    return web.Response(status=404)
  
  @server.PromptServer.instance.routes.get("/sdfx/workflow/list")
  @server.PromptServer.instance.routes.get("/sdfx/template/list")
  async def workflows_templates(request):
    list = []
    path = request.url.path

    if path == "/sdfx/workflow/list":
      working_dir = workflows_path
    elif path == "/sdfx/template/list":
      working_dir = templates_path
    
    for filename in os.listdir(working_dir):
      if filename.endswith(".json"):
        file_path = os.path.join(working_dir, filename)
        with open(file_path, 'r') as file:
          try:
            data = json.load(file)
            uid = data.get('uid')
            name = data.get('name')
            dateCreated = data.get('dateCreated')
            meta = data.get('meta')
            list.append({ 'uid': uid, 'name': name, 'meta': meta, 'dateCreated': dateCreated})

          except json.JSONDecodeError as e:
            return web.json_response({'error': 'invalid Json'}, status=400)
    return web.json_response(list)

  @server.PromptServer.instance.routes.get("/sdfx/workflow")
  @server.PromptServer.instance.routes.post("/sdfx/workflow")
  @server.PromptServer.instance.routes.put("/sdfx/workflow")
  @server.PromptServer.instance.routes.delete("/sdfx/workflow")
  @server.PromptServer.instance.routes.get("/sdfx/template")
  @server.PromptServer.instance.routes.post("/sdfx/template")
  @server.PromptServer.instance.routes.put("/sdfx/template")
  @server.PromptServer.instance.routes.delete("/sdfx/template")
  async def workflow_template(request):
    method = request.method
    path = request.url.path
    if method == 'GET':
      uid = request.query.get('uid')
      if not uid:
        return web.json_response({'error': 'uid missing'}, status=400)
      if path == "/sdfx/workflow":
        filename = os.path.join(workflows_path, f"{uid}.json")
      elif path == "/sdfx/template":
        filename = os.path.join(templates_path, f"{uid}.json")
      if not os.path.isfile(filename):
        return web.json_response({'error': 'file missing'}, status=404)
      with open(filename, 'r') as file:
        file_content = json.load(file)
      return web.json_response(file_content)
    elif method == 'POST' or method == 'PUT':
      try:
        data = await request.json()
        data['dateCreated'] = int(time.time())
        if(method == 'POST'):
          hash = str(uuid.uuid4())
        else:
          hash = request.query.get('uid')
          if not hash:
            return web.json_response({'error': 'uid missing'}, status=400)
        data['uid'] = hash
        if path == "/sdfx/workflow":
          filename = os.path.join(workflows_path, f"{hash}.json")
        elif path == "/sdfx/template":
          filename = os.path.join(templates_path, f"{hash}.json")
        with open(filename, 'w') as file:
          json.dump(data, file)
        response = {'uid': data.get('uid'), 'name': data.get('name'), 'meta': data.get('meta'), 'dateCreated': data.get('dateCreated')}
        return web.json_response(response)
      except json.JSONDecodeError:
        return web.json_response({'error': 'invalid Json'}, status=400)
    elif method == 'DELETE':
      uid = request.query.get('uid')
      if not uid:
        return web.json_response({'error': "Fatal: Parameter 'uid' missing"}, status=400)
      if path == "/sdfx/workflow":
        filename = os.path.join(workflows_path, f"{uid}.json")
      elif path == "/sdfx/template":
        filename = os.path.join(templates_path, f"{uid}.json")
      if os.path.isfile(filename):
        os.remove(filename)
      return web.json_response({'uid': uid, 'status':'deleted'})
    else:
      return web.Response(status=404)
    

  @server.PromptServer.instance.routes.get("/sdfx/input")
  @server.PromptServer.instance.routes.delete("/sdfx/input")
  @server.PromptServer.instance.routes.get("/sdfx/output")
  @server.PromptServer.instance.routes.delete("/sdfx/output")
  async def input_output(request):
    method = request.method
    path = request.url.path
    if method == 'GET':
      name = request.query.get('name')
      if not name:
        return web.json_response({'error': 'name missing'}, status=400)
      if path == '/sdfx/input':
        filename = os.path.join(folder_paths.get_input_directory(), name)
      elif path == '/sdfx/output':
        filename = os.path.join(folder_paths.get_output_directory(), name)
      if not os.path.isfile(filename):
        return web.json_response({'error': 'file missing'}, status=404)
      return web.json_response({'name': name})
    if method == 'DELETE':
      name = request.query.get('name')
      if not name:
        return web.json_response({'error': 'name missing'}, status=400)
      if path == '/sdfx/input':
        filename = os.path.join(folder_paths.get_input_directory(), name)
      elif path == '/sdfx/output':
        filename = os.path.join(folder_paths.get_output_directory(), name)
      if os.path.isfile(filename):
        os.remove(filename)
      return web.json_response({'name': name, 'status':'deleted'})
    
  @server.PromptServer.instance.routes.get("/sdfx/input/list")
  @server.PromptServer.instance.routes.get("/sdfx/output/list")
  async def inputs_outputs(request):
    list = []
    path = request.url.path

    if path == '/sdfx/input/list':
      working_dir = folder_paths.get_input_directory()
    elif path == '/sdfx/output/list':
      working_dir = folder_paths.get_output_directory()

    for filename in os.listdir(working_dir):
      if '.' in filename:
        list.append({'name': filename})
    return web.json_response(list)