import gradio as gr
import os
from pathlib import Path
import gc
import re
import shutil
from utils import set_token, get_download_file, list_uniq
from stkey import read_safetensors_key, read_safetensors_metadata, validate_keys, write_safetensors_key


TEMP_DIR = "."
KEYS_DIR = "keys"
KEYS_FILES = [f"{KEYS_DIR}/sdxl_keys.txt"]
DEFAULT_KEYS_FILE = f"{KEYS_DIR}/sdxl_keys.txt"


def update_keys_files():
    global KEYS_FILES
    files = []
    for file in Path(KEYS_DIR).glob("*.txt"):
        files.append(str(file))
    KEYS_FILES = files


update_keys_files()


def upload_keys_file(path: str):
    global KEYS_FILES
    newpath = str(Path(KEYS_DIR, Path(path).stem + ".txt"))
    if not Path(newpath).exists(): shutil.copy(str(Path(path)), newpath)
    update_keys_files()
    return gr.update(choices=KEYS_FILES)


def parse_urls(s):
    url_pattern = "https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+"
    try:
        urls = re.findall(url_pattern, s)
        return list(urls)
    except Exception:
        return []


def to_urls(l: list[str]):
    return "\n".join(l)


def uniq_urls(s):
    return to_urls(list_uniq(parse_urls(s)))


def download_file(dl_url, civitai_key, progress=gr.Progress(track_tqdm=True)):
    download_dir = TEMP_DIR
    progress(0, desc=f"Start downloading... {dl_url}")
    output_filename = get_download_file(download_dir, dl_url, civitai_key)
    return output_filename


def get_stkey(filename: str, is_validate: bool=True, rfile: str=KEYS_FILES[0], progress=gr.Progress(track_tqdm=True)):
    paths = []
    metadata = {}
    keys = []
    missing = []
    added = []
    try:
        progress(0, desc=f"Loading keys... {filename}")
        keys = read_safetensors_key(filename)
        if len(keys) == 0: raise Exception("No keys found.")
        progress(0.5, desc=f"Checking keys... {filename}")
        if write_safetensors_key(keys, str(Path(filename).stem + ".txt"), is_validate, rfile):
            paths.append(str(Path(filename).stem + ".txt"))
            paths.append(str(Path(filename).stem + "_missing.txt"))
            paths.append(str(Path(filename).stem + "_added.txt"))
        missing, added = validate_keys(keys, rfile)
        metadata = read_safetensors_metadata(filename)
    except Exception as e:
        print(f"Error: Failed check {filename}. {e}")
        gr.Warning(f"Error: Failed check {filename}. {e}")
    finally:
        Path(filename).unlink()
        return paths, metadata, keys, missing, added


def stkey_gr(dl_url: str, civitai_key: str, hf_token: str, urls: list[str], files: list[str], 
             is_validate=True, rfile=KEYS_FILES[0], progress=gr.Progress(track_tqdm=True)):
    if not hf_token: hf_token = os.environ.get("HF_TOKEN") # default huggingface token
    set_token(hf_token)
    if not civitai_key: civitai_key = os.environ.get("CIVITAI_API_KEY") # default Civitai API key
    dl_urls = parse_urls(dl_url)
    if not urls: urls = []
    if not files: files = []
    metadata = {}
    keys = []
    missing = []
    added = []
    for u in dl_urls:
        file = download_file(u, civitai_key)
        if not Path(file).exists() or not Path(file).is_file(): continue
        paths, metadata, keys, missing, added = get_stkey(file, is_validate, rfile)
        if len(paths) != 0: files.extend(paths)
    progress(1, desc="Processing...")
    gc.collect()
    return gr.update(value=urls, choices=urls), gr.update(value=files), gr.update(visible=False), metadata, keys, missing, added