#Utility functions for use in the nodes.

import pathlib
import hashlib
import datetime
import torch
import logging

import folder_paths
import comfy.utils

from .model_cache import cache
from .helpers_civitai import *

def str_to_bool(value):
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        value = value.lower()
        if value in {'true', '1', 'yes'}:
            return True
        if value in {'false', '0', 'no'}:
            return False
    raise ValueError(f"Cannot convert {value} to boolean.")

def bool_to_str(value):
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, str):
        value = value.lower()
        if value in {'true', '1', 'yes'}:
            return "true"
        if value in {'false', '0', 'no'}:
            return "false"
    raise ValueError(f"Cannot convert {value} to string representation of boolean.")

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

def get_file_sha256(path):
    print(f"Calculating hash for {path}")
    m = hashlib.sha256()
    with open(path, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            m.update(chunk)
    full_hash = m.digest().hex()
    print(f"Got full hash {full_hash}")
    result = full_hash[:10]
    print(f"Got hash {result}")
    return result


def get_files_in_dir(input_dirs=None, extensions=None):
    if extensions is None:
        extensions = ".*"
    if input_dirs is None:
        raise ValueError("input_dirs cannot be None")

    input_files = []
    # Check if input_dirs is a tuple or a list. If not, make it a list.
    if not isinstance(input_dirs, (list, tuple)):
        input_dirs = [input_dirs]

    for dir in input_dirs:
        if dir is None or dir == "":
            continue
        if pathlib.Path(dir).exists():
            file_list = pathlib.Path(dir).rglob("*")
            for file in file_list:
                if file.exists() and not file.is_dir():
                    file_path = ""
                    if file.suffix.lower() in extensions:
                        # Get path relative to the input directory
                        try:
                            file_path = str(file.relative_to(dir))
                        except ValueError:
                            file_path = str(file)
                    input_files.append(file_path)

    input_files = sorted(set(input_files))
    return input_files

def last_used(file_path):
    cache.load()

    if file_path in cache.hash:
        last_used = cache.by_path(file_path).get("lastUsed", None)
        if last_used is not None:
            return datetime.datetime.fromisoformat(last_used)
        else:
            return None
    else:
        return None

def days_since_last_used(file_path):
    was_last_used = last_used(file_path)
    if was_last_used is not None:
        now = datetime.datetime.now()
        delta = now - was_last_used
        return delta.days
    else:
        return 365

def get_file_modification_date(file_path):
    try:
        file_path = pathlib.Path(file_path)
        if file_path.exists():
            return datetime.datetime.fromtimestamp(file_path.stat().st_mtime)
        else:
            print(f"File {file_path} does not exist.")
            return datetime.datetime.now()
    except Exception as e:
        print(f"Error getting modification date for {file_path}: {e}")
        return datetime.datetime.now()

# Called *after* verifying the json pulled successfully.
# This function updates the cache with the information from the CivitAI json.
def update_cache_from_civitai_json(file_path, json_data, timestamp=True):
    the_files = json_data.get("files", [])
    hashes = {}

    if len(the_files) > 0:
        hashes = the_files[0].get("hashes", {})
    update_available = True

    latest_model = None
    if json_data.get("modelId", None) is not None:
        latest_model = get_latest_model_version(json_data["modelId"])
        if latest_model == json_data["id"] or latest_model is None:
            update_available = False
    
    if latest_model is None:
        latest_model = ""

    file_cache = cache.by_path(file_path)
    file_cache.update({
        'civitai': "True",
        'civitai_failed_count': 0,
        'model': json_data.get("model", {}),
        'name': json_data.get("name", ""),
        'baseModel': json_data.get("baseModel", ""),
        'id': json_data.get("id", ""),
        'modelId': json_data.get("modelId", ""),
        'update_available': update_available,
        'update_version_id': latest_model,
        'trainedWords': json_data.get("trainedWords", []),
        'downloadUrl': json_data.get("downloadUrl", ""),
        'hashes': hashes
    })
    if timestamp:
        print("Updating timestamp.")
        cache.update_last_used_by_path(file_path)

    print("Successfully pulled metadata.")

def update_cache_without_civitai_json(file_path, hash, timestamp=True):
    file_cache = cache.by_path(file_path)
    print(f"Unable to find on civitai.")
    file_cache['civitai'] = "False"
    file_cache['civitai_failed_count'] = file_cache.get('civitai_failed_count', 0) + 1
    file_cache['hash'] = hash
    cache.update_last_used_by_path(file_path)

def add_file_to_cache(file_path, hash=None):
    file_path = str(file_path)
    print(f"Adding {file_path} to cache.")
    if hash is None:
        hash = get_file_sha256(file_path)
    
    if file_path not in cache.hash:
        cache.hash[file_path] = hash
    
    if cache.info.get(hash, None) is None:
        cache.info[hash] = {
            'civitai': "False",
            'update_available': False,
            'update_version_id': "",
            'hash': hash,
            'lastUsed': datetime.datetime.now().isoformat()
        }
    
    print(f"Adding {file_path} to cache with hash {hash}.")
    return hash

def recheck_hash(file_path, hash):
    new_hash = get_file_sha256(file_path)
    if new_hash != hash:
        print(f"Hash mismatch. Using new hash.")
        if file_path in cache.hash:
            print(f"Updating cache for {file_path} with new hash {new_hash}.")
            if new_hash not in cache.info:
                if hash in cache.info:
                    cache.info[new_hash] = cache.info[hash]
        else:
            print(f"File {file_path} not in cache. Adding with new hash {new_hash}.")
            add_file_to_cache(file_path, new_hash)
        hash = new_hash
    return hash

def pull_metadata(file_paths, timestamp = True, force_all = False, pbar = None):
    pull_json = True
    metadata_days_recheck = 7
    
    cache.load()
    
    if isinstance(file_paths, str):
        file_paths = [file_paths]
    
    if not file_paths:
        print("No file paths provided.")
        return
    
    for file_path in file_paths:
        force = force_all
        #print(f"Processing file: {file_path}")
        #print(f"cache.hash: {cache.hash}")
        hash = cache.hash.get(str(file_path), None)
        #print(f"Current hash: {hash}")
        if hash is None:
            print(f"Hash not found in cache for {file_path}. Adding to cache.")
            hash = add_file_to_cache(file_path)

        file_cache = cache.by_path(file_path)

        last_used_date = datetime.datetime.fromisoformat(file_cache['lastUsed']) if 'lastUsed' in file_cache else None

        modified = get_file_modification_date(file_path)
        # If file was modified after last used, force metadata pull
        if last_used_date is not None and modified is not None and modified > last_used_date:
            print(f"File was modified after last used. Pulling metadata.")
            force = True
        
        # Only skip pull if not forced and civitai is True and recently pulled
        civitai_val = False
        try:
            civitai_val = str_to_bool(file_cache.get('civitai', False))
        except:
            civitai_val = False

        if not force and civitai_val == True:
            if days_since_last_used(file_path) <= metadata_days_recheck:
                print(f"Pulled metadata within the last {metadata_days_recheck} days. No pull needed.")
                pull_json = False

        # If force, recalculate hash before any API call
        if force:
            print(f"Force flag is set. Recalculating hash for {file_path}.")
            hash = recheck_hash(file_path, hash)

        if pull_json or force:
            print(f"Currently pulling metadata for {file_path}.")
            json = get_civitai_model_version_json_by_hash(hash)

            if 'error' in json:
                retried = False
                # Try fallback with modelId if available
                if 'modelId' in file_cache:
                    print(f"Using cached model id {file_cache.get('id', None)}")
                    json = get_civitai_model_version_json_by_id(file_cache['id'])
                    retried = True
                else:
                    print(f"No cached model id.")

                if 'error' in json:
                    if retried:
                        print(f"Error: {json['error']}")
                    print(f"Unable to find on civitai.")
                    file_cache['civitai'] = "False"
                    file_cache['civitai_failed_count'] = file_cache.get('civitai_failed_count', 0) + 1
                    update_cache_without_civitai_json(file_path, hash, timestamp=timestamp)
            
            if 'error' not in json:
                update_cache_from_civitai_json(file_path, json, timestamp=timestamp)
            else:
                retries = file_cache.get('civitai_failed_count', 0)
                retries += 1
                file_cache['civitai_failed_count'] = retries

        if timestamp:
            print("Updating timestamp.")
            cache.update_last_used_by_path(file_path)

        cache.hash[file_path] = hash
        cache.info[hash] = file_cache
        if pbar is not None:
            pbar.update(1)
    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.hash[lora_path]

def model_scan(the_path, force = False):
    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))
    model_list = [str(x) for x in model_list]
    print(f"Scanning {len(model_list)} models for metadata.")
    pbar = comfy.utils.ProgressBar(len(model_list))
    pull_metadata(model_list, force_all=force, timestamp=False, pbar=pbar)

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.hash.keys():
            continue
        
        if 'lastUsed' not in cache.by_path(model_path):
            continue

        last = cache.by_path(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 clean_keywords(keywords):
    keywords = set(filter(None, (x.strip() for x in keywords)))
    return ', '.join(keywords)

def clean_text(text):
    ret = ' '.join(filter(None, (x.strip() for x in text.split())))
    ret = ', '.join(filter(None, (x.strip() for x in ret.split(','))))
    ret = '\n'.join(filter(None, (x.strip() for x in ret.split('\n'))))
    
    
    # Strip whitespace from the start and end of text in parentheses
    ret = ' ('.join(part.strip() for part in ret.split('('))
    ret = ')'.join(part.strip() for part in ret.split(')'))
    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]]


def get_save_file_path(filename_prefix: str = "text", filename_ext: str = "txt") -> str:
    """
    Generate a safe file path for saving files with automatic counter increment.

    Args:
        filename_prefix: Base filename, can include date/time variables like %year%, %month%, etc.
        filename_ext: File extension (without dot)

    Returns:
        Complete file path including directory and filename with counter
    """

    def _extract_counter_from_filename(filename: str) -> tuple[int, str]:
        """Extract counter from existing filename to determine next counter value."""
        base_name = pathlib.Path(filename_prefix).name
        prefix_len = len(base_name)

        if len(filename) <= prefix_len + 1:
            return 0, filename[:prefix_len + 1]

        prefix = filename[:prefix_len + 1]
        try:
            # Remove file extension first, then extract counter
            filename_no_ext = pathlib.Path(filename).stem
            counter_part = filename_no_ext[prefix_len + 1:]
            digits = int(counter_part)
        except (ValueError, IndexError):
            digits = 0
        return digits, prefix

    def _replace_date_variables(text: str) -> str:
        """Replace date/time variables in the filename prefix."""
        now = datetime.datetime.now()
        replacements = {
            "%year%": str(now.year),
            "%month%": str(now.month).zfill(2),
            "%day%": str(now.day).zfill(2),
            "%hour%": str(now.hour).zfill(2),
            "%minute%": str(now.minute).zfill(2),
            "%second%": str(now.second).zfill(2),
        }

        for placeholder, value in replacements.items():
            text = text.replace(placeholder, value)
        return text

    # Get output directory and process filename prefix
    output_dir = folder_paths.get_output_directory()
    
    if "%" in filename_prefix:
        filename_prefix = _replace_date_variables(filename_prefix)

    # Parse the filename prefix path
    filename_prefix_path = pathlib.Path(filename_prefix)
    subfolder = filename_prefix_path.parent
    base_filename = filename_prefix_path.name

    # Construct full output path
    output_path = pathlib.Path(output_dir)
    full_output_folder = output_path / subfolder

    # Security check: ensure we're not saving outside the output directory
    try:
        full_output_folder.resolve().relative_to(output_path.resolve())
    except ValueError:
        error_msg = (
            "ERROR: Saving outside the output folder is not allowed.\n"
            f"  Target folder: {full_output_folder.resolve()}\n"
            f"  Output directory: {output_path.resolve()}"
        )
        logging.error(error_msg)
        raise ValueError(error_msg)

    # Ensure output directory exists
    full_output_folder.mkdir(parents=True, exist_ok=True)

    # Find the next available counter
    counter = 1
    try:
        existing_files = [f.name for f in full_output_folder.iterdir() if f.is_file()]
        matching_counters = []

        for file in existing_files:
            digits, prefix = _extract_counter_from_filename(file)
            # Check if this file matches our pattern (same base name and ends with underscore)
            if (prefix[:-1].lower() == base_filename.lower() and 
                len(prefix) > 0 and prefix[-1] == "_"):
                matching_counters.append(digits)

        if matching_counters:
            counter = max(matching_counters) + 1
            
    except Exception as e:
        logging.warning(f"Error finding existing files, using counter=1: {e}")
        counter = 1

    # Generate final filename
    final_filename = f"{base_filename}_{counter:05d}.{filename_ext}"

    return str(full_output_folder / final_filename)

def unwrap_tuple(value):
    """Unwrap single-item tuples to their contained value."""
    return value[0] if isinstance(value, tuple) and len(value) == 1 else value