#Utility functions for use in the nodes.

import pathlib
import hashlib
import requests
import time
import datetime
import numpy as np
import torch
import json
from PIL import Image, ImageOps

import folder_paths
import comfy.utils

from . import cache

def name_from_path(path):
    return pathlib.Path(path).name

def get_civitai_model_version_json(hash):
    try:
        r = requests.get("https://civitai.com/api/v1/model-versions/by-hash/" + hash)
        r.raise_for_status()
    except HTTPError as http_err:
        print(f"HTTP error occurred: {http_err}")
        return {"error": "HTTP error occurred: " + http_err}
    except Exception as err:
        print(f"Other error occurred: {err}")
        return {"error": "Other error occurred: " + err}
    else:
        print("Retrieved json from civitai.")
        return r.json()

    return r.json()

def get_civitai_model_json(modelId):
    try:
        r = requests.get("https://civitai.com/api/v1/models/" + str(modelId))
        r.raise_for_status()
    except HTTPError as http_err:
        print(f"HTTP error occurred: {http_err}")
        return {"error": "HTTP error occurred: " + http_err}
    except Exception as err:
        print(f"Other error occurred: {err}")
        return {"error": "Other error occurred: " + err}
    else:
        print("Retrieved json from civitai.")
        return r.json()

    return r.json()

def get_file_sha256(path):
    print(f"Calculating hash for {path}")
    m = hashlib.sha256()
    
    with open(path, 'rb') as f:
        m.update(f.read())

    result = str(m.digest().hex()[:10])
    print(f"Got hash {result}")
    return result

def pull_metadata(file_path, timestamp = False):
    cache.load()
    
    print(f"Pull metadata for {file_path}.")
    hash = cache.data.get(file_path, {}).get("hash", "")

    if not hash:
        cache.data[file_path] = {"hash": get_file_sha256(file_path)}
        hash = cache.data[file_path]["hash"]
    else:
        time.sleep(2)

    try:
        pull_json = True
        file_cache = cache.data.get(file_path, {})
        
        if 'lastUsed' in file_cache and 'civitai' in file_cache:
            last_used = datetime.datetime.fromisoformat(file_cache['lastUsed'])
            if (datetime.datetime.now() - last_used).days == 0:
                print(f"Pulled earlier today. No pull needed. Update timestamp = {timestamp}")
                pull_json = False

        if pull_json:
            json = get_civitai_model_version_json(hash)
            if 'error' in json:
                print(f"Error: {json['error']}")
                file_cache['civitai'] = file_cache.get('model', "False")
            else:
                file_cache.update({
                    'civitai': "True",
                    'model': json["model"],
                    'name': json["name"],
                    'baseModel': json["baseModel"],
                    'id': json["id"],
                    'modelId': json["modelId"],
                    'trainedWords': json["trainedWords"],
                    'downloadUrl': json["downloadUrl"]
                })
                print("Successfully pulled metadata.")
    except Exception as e:
        print(f"Failed to pull metadata for {file_path} with hash {hash}: {e}")
        file_cache['civitai'] = file_cache.get('civitai', "False")

    if timestamp:
        print("Updating timestamp.")
        file_cache['lastUsed'] = datetime.datetime.now().isoformat()

    cache.data[file_path] = file_cache
    cache.save()

def lora_to_string(lora_name, model_weight, clip_weight):
    lora_string = ' <lora:' + str(pathlib.Path(lora_name).name) + ":" + str(model_weight) +  ">" #  + ":" + str(clip_weight)
        
    return lora_string

def lora_to_prompt(lora_stack = None):
    lora_info = ''
    if lora_stack is None:
        return ""
    else:
        for lora in lora_stack:
            lora_info += lora_to_string(lora[0], lora[1], lora[2])
    return lora_info

def get_lora_hash(lora_name):
    lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
    pull_metadata(lora_path)

    return cache.data[lora_path]["hash"]

def get_model_info(lora_path, weight = None):
    ret = {}
    try:
        ret["type"] = cache.data[lora_path]["model"]["type"]
        if (ret["type"] == "LORA") and (weight is not None):
            ret["weight"] = weight
        ret["modelVersionId"] = cache.data[lora_path]["id"]
        ret["modelName"] = cache.data[lora_path]["model"]["name"]
        ret["modelVersionName"] = cache.data[lora_path]["name"]
    except:
        ret = {}
    return ret

def get_latest_model_version(modelId):
    json = get_civitai_model_json(modelId)
    if 'error' in json:
        return json['error']

    latest_model = None
    model_date = None
    for model in json["modelVersions"]:
        if model_date is None or (datetime.datetime.fromisoformat(model['createdAt']) > model_date and model['status'] == "Published" and model['availability'] == "Public"):
            model_date = datetime.datetime.fromisoformat(model['createdAt'])
            latest_model = model["id"]

    return latest_model

def model_scan(the_path):
    the_paths = the_path

    print(f"the_paths: {the_paths}")

    model_list = []
    for dir in the_paths:
        print(f"dir: {dir}")
        result = list(p.resolve() for p in pathlib.Path(dir).glob("**/*") if p.suffix in {".safetensors", ".ckpt"})
        model_list.extend(result)

    model_list = list(set(model_list))
    print(f"There are {len(model_list)} files.")
    pbar = comfy.utils.ProgressBar(len(model_list))
    for the_model in model_list:
        pbar.update(1)
        pull_metadata(str(the_model))

def pull_lora_image_urls(hash, nsfw):
    json = get_civitai_model_version_json(hash)
    img_list = []
    for pic in json['images']:
        if pic['nsfwLevel'] > 1:
            if nsfw == True:
                img_list.append(pic['url'])
        else:
            img_list.append(pic['url'])
    return img_list

def url_to_torch_image(url):
    img = Image.open(requests.get(url, stream=True).raw)
    img = ImageOps.exif_transpose(img)
    img = np.array(img.convert("RGB")).astype(np.float32) / 255.0
    return (torch.from_numpy(img)[None,])

def blank_image():
    img = Image.new('RGB', (1024, 1024))
    img = ImageOps.exif_transpose(img)
    img = np.array(img.convert("RGB")).astype(np.float32) / 255.0
    return (torch.from_numpy(img)[None,])

def get_recently_used_models(model_type):
        model_list = list()
        full_model_list = folder_paths.get_filename_list(model_type)
        for item in full_model_list:
            model_path = folder_paths.get_full_path_or_raise(model_type, item)
            if model_path not in cache.data.keys():
                continue
            
            if 'lastUsed' not in cache.data[model_path]:
                continue
            
            last = cache.data[model_path]['lastUsed']
            last_used = datetime.datetime.fromisoformat(last)
            #print(f"{model_path} - last: {last} last_used: {last_used}")
            if (datetime.datetime.now() - last_used).days <= 7:
                model_list.append(item)
        return model_list

def civitai_sampler_name(sampler_name, scheduler_name):
    comfy_to_auto = {
        'ddim': 'DDIM',
        'dpm_2': 'DPM2',
        'dpm_2_ancestral': 'DPM2 a',
        'dpmpp_2s_ancestral': 'DPM++ 2S a',
        'dpmpp_2m': 'DPM++ 2M',
        'dpmpp_sde': 'DPM++ SDE',
        'dpmpp_2m_sde': 'DPM++ 2M SDE',
        'dpmpp_2m_sde_gpu': 'DPM++ 2M SDE',
        'dpmpp_3m_sde': 'DPM++ 3M SDE',
        'dpmpp_3m_sde_gpu': 'DPM++ 3M SDE',
        'dpm_fast': 'DPM fast',
        'dpm_adaptive': 'DPM adaptive',
        'euler_ancestral': 'Euler a',
        'euler': 'Euler',
        'heun': 'Heun',
        'lcm': 'LCM',
        'lms': 'LMS',
        'plms': 'PLMS',
        'uni_pc': 'UniPC',
        'uni_pc_bh2': 'UniPC'
    }
    result = comfy_to_auto.get(sampler_name, sampler_name)

    if (scheduler_name == "karras"):
        result += " Karras"
    elif (scheduler_name == "exponential"):
        result += " Exponential"

    return result

def clean_keywords(keywords):
    keywords = list(set(keywords))
    keywords = [x for x in keywords if x != '']
    keywords = [x for x in keywords if x != None]
    keywords = [x for x in keywords if x != ' ']
    
    ret = ' '.join(", ".join(keywords).split('\n'))
    return ret

def clean_text(text):
    ret_list = [x for x in text.split(" ") if x.strip()]
    ret = " ".join(ret_list)

    ret_list = [x for x in ret.split(",") if x.strip()]
    ret = ", ".join([x.strip(" ") for x in ret_list])

    ret_list = [x for x in ret.split("\n") if x.strip()]
    ret = "\n".join([x.strip(" ") for x in ret_list])
    return (ret)

def condition_text(clip, text = None):
    zero_text = text is None
    text = text or ""
    
    tokens = clip.tokenize(text)
    output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
    cond = output.pop("cond")

    if zero_text:
        pooled_output = output.get("pooled_output")
        if pooled_output is not None:
            output["pooled_output"] = torch.zeros_like(pooled_output)
        return [[torch.zeros_like(cond), output]]

    return [[cond, output]]