import json
import urllib3
import urllib.request
import os
import cgi
from pathlib import Path
import argparse
import string
#import getpass
#print("Env thinks the user is [%s]" % (os.getlogin()))
#print("Effective user is [%s]" % (getpass.getuser()))

import erisscan


def filter_str(s: str) -> str:
    return "".join(filter(lambda x: x in string.printable, s))


def downloadcivitai(modelid: str, type: str, image: bool) -> bool:
    
    match type:
        case "lora":
            dlpath = f"./models/Lora/"
        case "sd":
            dlpath = f"./models/Stable-diffusion/"
        case "emb":
            dlpath = f"./embeddings/"

    http = urllib3.PoolManager()
    urlid = f"https://civitai.com/api/v1/models/{modelid}"
    resp = http.request('GET', urlid)
    modelarray = json.loads(resp.data.decode('utf-8'))
    modelsha256 = modelarray.get('modelVersions')[0].get('files')[0].get('hashes').get("SHA256")
    urlhash = f"https://civitai.com/api/v1/model-versions/by-hash/{modelsha256}"
    modelimageurl = modelarray.get('modelVersions')[0].get('images')[0].get('url')
    modelfilename = modelarray.get('modelVersions')[0].get('files')[0].get('name')
    #modelfilesize = int(modelarray.get('modelVersions')[0].get('files')[0].get('sizeKB'))
    modelname = Path(modelfilename).stem
    modelfileurl = modelarray.get('modelVersions')[0].get('files')[0].get('downloadUrl')
    #modeldisplayname = modelarray.get('name')
    #trainedwords = modelarray.get('modelVersions')[0].get('trainedWords')

    modelimagename = [f"{modelname}.preview.jpg", f"{modelname}.preview.jpeg", f"{modelname}.preview.png"]
    #modelimagename = ""
    modelinfofile = f"{modelname}.civitai.info"
    
    
    def downloadmodel(url: str, dlpath: str) -> bool:
        print(f"Download {modelfilename} to {dlpath}")
        dlrequest = http.request('GET', url, preload_content=False)
        filesize = dlrequest.getheader('content-length')
        contentdis = dlrequest.getheader('Content-Disposition')
        try:
            type, filename = cgi.parse_header(contentdis)
            #print("Found file", dict(filename).get('filename'))
            filename = dict(filename).get('filename')
            #filename = filename.encode('utf-8')
            #filename = filter_str(filename)
        except Exception as e:
            print("Cannot get filename", e)
        if filesize:
            filesize = int(filesize)
            blocksize = max(4096, filesize//100)
            print(f"Filesize: {filesize//1024} kB | {filesize//1024//1024} MB")
        else:
            blocksize = 10000
        #with open(f"{dlpath}{filename}", 'wb') as out:
        with open(Path(dlpath) / filename, 'wb') as out:
            size = 0
            while True:
                data = dlrequest.read(blocksize)
                if not data:
                    break
                size += len(data)
                out.write(data)
                if filesize:
                    Percent = int((size / filesize)*100)
                    if Percent == 100:
                        print(f"Download: {Percent}%")
                    else:
                        print(f"Download: {Percent}%", end='\r')
        dlrequest.release_conn()
        
        
        
    def downloadimage(url: str, dlpath: str) -> bool:
        print(f"Download Image for {modelfilename}")
        imagerequest = http.request('GET', url, preload_content=False)
        #path = f"{os.getcwd()}/models/Loras/"
        filesize = imagerequest.getheader('content-length')
        #contentdis = imagerequest.getheader('Content-Disposition')
        #type, filename = cgi.parse_header(contentdis)
        filename = f"{Path(modelfilename).stem}.preview{Path(filename).suffix}"
        global modelimagename
        modelimagename = filename
        #print(filename['filename'])
        if filesize:
            filesize = int(filesize)
            blocksize = max(4096, filesize//100)
            print(f"Filesize: {filesize//1024} kB | {filesize//1024//1024} MB")
        else:
            blocksize = 1000
        with open(f"{dlpath}{filename}", 'wb') as out:
            size = 0
            while True:
                data = imagerequest.read(blocksize)
                if not data:
                    break
                size += len(data)
                out.write(data)
                if filesize:
                    Percent = int((size / filesize)*100)
                    if Percent == 100:
                        print(f"Download: {Percent}%")
                    else:
                        print(f"Download: {Percent}%", end='\r')
        imagerequest.release_conn()
        
    def downloadinfo(url: str, dlfullpath: str) -> bool:
        print(f"Download Infofile for {modelfilename}")
        resp = http.request('GET', url)
        infofile = open(f"{dlfullpath}", "w")
        infofile.write(resp.data.decode('utf-8'))
        infofile.close()
    
    # download Model file
    try:
        if not os.path.exists(f"{dlpath}{modelfilename}"):
            downloadmodel(modelfileurl, dlpath)
        else:
            print(f"{modelfilename} already exist, skip download")
    except Exception as e:
        print(e)
    finally:
        if not os.path.exists(f"{dlpath}{modelfilename}"):
            print(f"{modelfilename} could not be saved, abort.")
            return False
    # download Preview image, if not disabled   
    if image:
        try:
            if os.path.exists(f"{dlpath}{modelimagename[0]}") == False and os.path.exists(f"{dlpath}{modelimagename[1]}") == False and os.path.exists(f"{dlpath}{modelimagename[2]}") == False:
                downloadimage(modelimageurl, dlpath)
            else:
                print(f"Preview Image for {modelfilename} already exist, skip download")
                
        except Exception as e:
            print(e)
        finally:
            if isinstance(modelfilename, str):
                if not os.path.exists(f"{dlpath}{modelimagename}"):
                    print(f"{modelfilename} could not be saved, continue.")
    # save infofile
    try:
        if not os.path.exists(f"{dlpath}{modelinfofile}"):
            downloadinfo(urlhash, f"{dlpath}{modelinfofile}")
        else:
            print(f"{modelinfofile} already exist, skip creation")
    except Exception as e:
        print(e)
    finally:
        if not os.path.exists(f"{dlpath}{modelinfofile}"):
            print(f"{modelinfofile} could not be created, abort.")
            return False
    return True
    
# ______________________________________________________________________  

def downloadsideload(name: str, modeltype: str) -> bool:
    sideloadserver = f"https://eris.hitmare.me/eris/{modeltype}/"
    
    match modeltype:
        case "lora":
            dlpath = "./models/lora/"
        case "sd":
            dlpath = "./models/stable-diffusion/"
        case "emb":
            dlpath = "./embeddings/"
    
    http = urllib3.PoolManager()
    url = f"{sideloadserver}{name}"
    try:
        print(f"Download {name}")
        dlrequest = http.request('GET', url, preload_content=False)
        
        if not int(dlrequest.status) == 200:
            print(f"File {name} not found")
            return False
        
        filesize = dlrequest.getheader('content-length')
        #print(filesize)
        #contentdis = dlrequest.getheader('Content-Disposition')
        #type, filename = cgi.parse_header(contentdis)
        filename = name
        if filesize:
            filesize = int(filesize)
            blocksize = max(4096, filesize//100)
            print(f"Filesize: {filesize//1024} kB | {filesize//1024//1024} MB")
        else:
            blocksize = 10000
        with open(f"{dlpath}{filename}", 'wb') as out:
            size = 0
            while True:
                data = dlrequest.read(blocksize)
                if not data:
                    break
                size += len(data)
                out.write(data)
                if filesize:
                    Percent = int((size / filesize)*100)
                    if Percent == 100:
                        print(f"Download: {Percent}%")
                    else:
                        print(f"Download: {Percent}%", end='\r')
            print(f"Downloaded {dlpath}{filename}")
        dlrequest.release_conn()
        return True
    except Exception as e:
        print(name, "not found on Sideloadserver", " | Python Error:" ,e)
        return False
        
    

# ______________________________________________________________________  
# check current status of online and local lora
def checkcivitai(locallist: list, erislist: list) -> tuple:
    #print("Check Civitai")
    #url = "https://eris.hitmare.me/lora.txt"
    erismodels = erislist
    localmodels = locallist
    missingmodels = set()
    additionalmodels = set()
    savedmodels = set()        
    # check online list against local list
    for loraid in erismodels:
        if not loraid in localmodels:
            missingmodels.add(loraid)
        else:
            savedmodels.add(loraid)
    
    # check local list against online list
    for loraid in localmodels:
        if not loraid in erismodels:
            additionalmodels.add(loraid)
            
    return tuple((missingmodels, additionalmodels, savedmodels))


def checksideload(locallist: list, erislist: list) -> tuple:
    #print("Check Sideload")
    #url = "https://eris.hitmare.me/lora.txt"
    erismodels = erislist
    localmodels = locallist
    missingmodels = set()
    additionalmodels = set()
    savedmodels = set()        
    # check online list against local list
    for loraid in erismodels:
        if not loraid in localmodels:
            missingmodels.add(loraid)
        else:
            savedmodels.add(loraid)
    
    # check local list against online list
    for loraid in localmodels:
        if not loraid in erismodels:
            additionalmodels.add(loraid)
            
    return tuple((missingmodels, additionalmodels, savedmodels))

def deleteFiles(nameid: str, modeltype: str, localfile: bool = False) -> bool:
    match modeltype:
        case "lora":
            dlpath = "./models/lora/"
        case "sd":
            dlpath = "./models/stable-diffusion/"
        case "emb":
            dlpath = "./embeddings/"
    
    try: 
        if localfile:
            modelfilename = nameid
        else:
            http = urllib3.PoolManager()
            urlid = f"https://civitai.com/api/v1/models/{nameid}"
            resp = http.request('GET', urlid)
            modelarray = json.loads(resp.data.decode('utf-8'))
            modelfilename = modelarray.get('modelVersions')[0].get('files')[0].get('name')
        
        modelfiles = os.listdir(dlpath)
        for modelfile in modelfiles:
            if Path(modelfile).stem == Path(modelfilename).stem:
                os.remove(f"{dlpath}{modelfile}")
        return True
    except Exception as e:
        print(f"Could not delete {nameid}. Error: {e}")
        return False

parser = argparse.ArgumentParser(
                    prog='Eris File Updater',
                    description='Updates Lora, Embeddings and Stable-diffusion Models',
                    epilog='Thanks for being a part of the Eris Bot ^-^')

parser.add_argument('-si', '--skipimage', help=r"Skips the download of the preview images, Applies to all Model Types", action="store_true")
parser.add_argument('-sl', '--skiplora', help=r"Skips check of lora models", action="store_true")
parser.add_argument('-ss', '--skipsdmodel', help=r"Skips check of stable-diffusion models", action="store_true")
parser.add_argument('-se', '--skipemb', help=r"Skips check of embedding models", action="store_true")
parser.add_argument('-c', '--cleanup', help=r"Deletes any models that are locally but not on the eris model list. Affected by Skip Triggers", action="store_true")

args = parser.parse_args()   

if not os.path.exists("./eris.json"):
    # start scan python or exit
    #print("Please run the erisscan bat/sh file first")
    #exit()
    print("No eris.json file found. Starting scan ...")
    erisscan.scanall()
    
    
# get the eris list from the web and local
http = urllib3.PoolManager()
erisurl = "https://eris.hitmare.me/eris.json"
resp = http.request('GET', erisurl)
erisremote = json.loads(resp.data.decode('utf-8'))
erislocal = json.load(open("./eris.json", "r"))

if args.skipimage:
    dlimage = False
else:
    dlimage = True
    
if args.skiplora:
    print("Skip Lora Model check")
else:
    lorastate = checkcivitai(erislocal.get('lora')['civitai'], erisremote.get('lora')['civitai'])
    if len(lorastate[0]) > 0 :
        print(f"Missing lora ids: {lorastate[0]}")
        for lora in lorastate[0]:
            if downloadcivitai(lora, 'lora', dlimage):
                erislocal.get('lora')['civitai'].append(lora)
    lorasideloadstate = checksideload(erislocal.get('lora')['sideload'], erisremote.get('lora')['sideload'])
    if len(lorasideloadstate[0]) > 0 :
        print(f"Missing lora files: {lorasideloadstate[0]}")
        for lora in lorasideloadstate[0]:
            if downloadsideload(lora, 'lora'):
                erislocal.get('lora')['sideload'].append(lora)
    if args.cleanup:
        if len(lorastate[1]) > 0:
            print("Additional CivitAI lora ids:", lorastate[1])
            for dellora in lorastate[1]:
                if deleteFiles(dellora, 'lora'):
                    erislocal.get('lora')['civitai'].remove(dellora)
        if len(lorasideloadstate[1]) > 0:
            print("Additional CivitAI lora files:", lorastate[1])
            for delslora in lorasideloadstate[1]:
                if deleteFiles(delslora, 'lora', True):
                    erislocal.get('lora')['sideload'].remove(delslora)
    #print(len(lorastate[1]))
if args.skipsdmodel:
    print("Skip Stable-Diffusion Model check")
else:
    sdstate = checkcivitai(erislocal.get('sd')['civitai'], erisremote.get('sd')['civitai'])
    if len(sdstate[0]) > 0 :
        print(f"Missing CivitAI stable-diffusion ids: {sdstate[0]}")
        for sd in sdstate[0]:
            if downloadcivitai(sd, 'sd', dlimage):
                erislocal.get('sd')['civitai'].append(sd)
    sdsideloadstate = checksideload(erislocal.get('sd')['sideload'], erisremote.get('sd')['sideload'])
    if len(sdsideloadstate[0]) > 0 :
        print(f"Missing Sideload stable-diffusion files: {sdsideloadstate[0]}")
        for sd in sdsideloadstate[0]:
            if downloadsideload(sd, 'sd'):
                erislocal.get('sd')['sideload'].append(sd)
    if args.cleanup:
        if len(sdstate[1]) > 0:
            print("Additional CivitAI stable-diffusion ids:", sdstate[1])
            for delsd in sdstate[1]:
                if deleteFiles(delsd, 'lora'):
                    erislocal.get('sd')['civitai'].remove(delsd)
        if len(sdsideloadstate[1]) > 0:
            print("Additional Sideload stable-diffusion files:", sdstate[1])
            for delssd in sdsideloadstate[1]:
                if deleteFiles(delssd, 'lora', True):
                    erislocal.get('sd')['sideload'].remove(delssd)
                
if args.skipemb:
    print("Skip Embedding Model check")
else:
    embstate = checkcivitai(erislocal.get('emb')['civitai'], erisremote.get('emb')['civitai'])
    if len(embstate[0]) > 0 :
        print(f"Missing embedding ids: {embstate[0]}")
        for emb in embstate[0]:
            if downloadcivitai(emb, 'emb', dlimage):
                erislocal.get('emb')['civitai'].append(emb)
    embsideloadstate = checksideload(erislocal.get('emb')['sideload'], erisremote.get('emb')['sideload'])
    if len(embsideloadstate[0]) > 0 :
        print(f"Missing embedding files: {embsideloadstate[0]}")
        for emb in embsideloadstate[0]:
            if downloadsideload(emb, 'emb'):
                erislocal.get('emb')['sideload'].append(emb)
    if args.cleanup:
        if len(embstate[1]) > 0:
            print("Additional CivitAI embedding ids:", embstate[1])
            for delemb in embstate[1]:
                if deleteFiles(delemb, 'lora'):
                    erislocal.get('emb')['civitai'].remove(delemb)
        if len(embsideloadstate[1]) > 0:
            print("Additional Sideload embedding files:", embsideloadstate[1])
            for delsemb in embsideloadstate[1]:
                if deleteFiles(delsemb, 'lora', True):
                    erislocal.get('emb')['sideload'].remove(delsemb)

erisjsonfile = open("./eris.json", "w")
erisjsonfile.write(json.dumps(erislocal, indent=2))
erisjsonfile.close()