import os
import json
import hashlib
import urllib3
from pathlib import Path


def scanfolder(path: str) -> list:
    civitailist = []
    sideloaddict = []
    modelfolder = path
    modelfiles = os.listdir(modelfolder)
    for modelfile in modelfiles:
        modelname = Path(modelfile).stem
        sideload = False
        if modelfile.endswith('.safetensors') or modelfile.endswith('.pt'):
            #print(f"Found '{lorafile}' - Searching civitai.info")
            if os.path.exists(f"{modelfolder}{modelname}.civitai.info"):
                civitaiinfo = open(f"{modelfolder}{modelname}.civitai.info", "r", encoding='utf-8')
                civitaiinfocontent = civitaiinfo.read()
                civitaiinfo.close()
                if civitaiinfocontent[:2] == "{}":
                    os.remove(f"{modelfolder}{modelname}.civitai.info")
                    sideload = True
                else:
                    with open(f"{modelfolder}{modelname}.civitai.info", "r") as f:
                        #fstring = "dd"
                        print("Reading", modelfolder, modelfile)
                        modelinfoarray = json.load(f)
                        #print(loraarray)
                        try:
                            modelsha256 = modelinfoarray.get('files')[0].get('hashes').get('SHA256')
                            modelid = modelinfoarray.get('modelId')
                            civitailist.append(modelid)
                            #print(lorasha256, " " , loramodelid)
                        except Exception as e:
                            print("[Model Step local civitai.info parse] ", path, modelfile , " | ",e) 
                            #print("Content:", civitaiinfocontent[:2])
            else:
                sideload = True
            if sideload:
                print(f"{modelfolder}{modelname}.civitai.info not found. Getting SHA256 Hash and fetch data from civitai API")
                with open(f"{modelfolder}{modelfile}", "rb") as f:
                    bytes = f.read() # read entire file as bytes
                    modelsha256 = hashlib.sha256(bytes).hexdigest();
                    #print(readable_hash)
                    http = urllib3.PoolManager()
                    url = f"https://civitai.com/api/v1/model-versions/by-hash/{modelsha256}"
                    resp = http.request('GET', url)
                    modelinfoarray = json.loads(resp.data.decode('utf-8'))
                    try:
                        if modelinfoarray.get('error') == "Model not found":
                                print(f"Model {modelfile} is not a CivitAi Model")
                                sideloaddict.append(modelfile)
                                if os.path.exists(f"{modelfolder}{modelname}.preview.jpg"):
                                    sideloaddict.append(f"{modelname}.preview.jpg")
                                if os.path.exists(f"{modelfolder}{modelname}.preview.jpeg"):
                                    sideloaddict.append(f"{modelname}.preview.jpeg")
                                if os.path.exists(f"{modelfolder}{modelname}.preview.png"):
                                    sideloaddict.append(f"{modelname}.preview.png")
                        else:
                            print(f"Model {modelfile} found on CivitAi")
                            modelid = modelinfoarray.get('modelId')
                            civitailist.append(modelid)
                            print(f"Creating {modelfolder}{modelname}.civitai.info for future use")
                            modelinfofile = open(f"{modelfolder}{modelname}.civitai.info", "w")
                            modelinfofile.write(resp.data.decode('utf-8'))
                            modelinfofile.close()
                            
                    except Exception as e:
                        print("[Model Step SHA256 civitai.info parse] ", path, modelfile , " | ",e)
            #if os.path.exists(f"{modelfolder}{modelname}.civitai.info"):
                #print("Civitai Info file found.")
                # do json stuff
                
            #else:
                
    return list([civitailist, sideloaddict])
            

def scanall():
    erislist = {}
    erislist['lora'] = {}
    result = scanfolder("./models/lora/")
    erislist['lora']['civitai'] = result[0]
    erislist['lora']['sideload'] = result[1]
    
    erislist['sd'] = {}
    result = scanfolder("./models/stable-diffusion/")
    erislist['sd']['civitai'] = result[0]
    erislist['sd']['sideload'] = result[1]
    
    erislist['emb'] = {}
    result = scanfolder("./embeddings/")
    erislist['emb']['civitai'] = result[0]
    erislist['emb']['sideload'] = result[1]

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


if __name__ == "__main__":
    scanall()