import numpy as np
import mikeio
import os
import shutil
import zipfile

def uz2u10(Uz, z, mode, z0=0.001, power=0.1):
    if mode == "log":
        kappa = 0.4  # von Karman constant
        Uf = Uz * kappa / np.log(z / z0)
        U10 = Uf * (1 / kappa) * np.log(10 / z0)
    elif mode == "power":
        U10 = Uz * (10 / z) ** power
    else:
        raise ValueError("Mode must be 'log' or 'power'")
    return U10


def find_dfs2_index(ds, x, y):
    X = ds.geometry.x
    Y = ds.geometry.y
    dx = ds.geometry.dx
    dy = ds.geometry.dy
    j = np.argmin(np.abs(X - x))
    i = np.argmin(np.abs(Y - y))
    # Check if the point is within the grid cell
    if (X[j] - dx / 2 <= x <= X[j] + dx / 2) and (Y[i] - dy / 2 <= y <= Y[i] + dy / 2):
        return i, j
    else:
        raise ValueError("Point is outside the grid cell")
    

##### Bundle a MIKE setup file and its dependencies into a zip file #####
def _find_inputs(data, path=None, results=None):
    if path is None:
        path = []
    if results is None:
        results = []
    if isinstance(data, dict):
        for key, value in data.items():
            _find_inputs(value, path + [key], results)
    else:
        if isinstance(data, str) and data.startswith("|") and data.endswith("|"):
            results.append((path, data))
    return results

def _get_fullname(fname):
    if ".." in fname:
        return os.path.abspath(fname)
    return fname

def _existing(path, fname):
    if os.path.exists(os.path.join(path, fname)):
        basename = os.path.basename(fname)
        filename = os.path.splitext(basename)[0]
        extension = os.path.splitext(basename)[1]
        counter = 1
        while True:
            new_fname = f"{filename}_{counter}{extension}"
            if not os.path.exists(os.path.join(path, new_fname)):
                return new_fname
            counter += 1
    return fname

def _set_value(data, path, new_value):
    current = data
    for key in path[:-1]:
        try:
            current = current[key]
        except (KeyError, IndexError, TypeError):
            raise ValueError(f"Invalid path at {key}")
    last_key = path[-1]
    try:
        current[last_key] = new_value
    except (KeyError, IndexError, TypeError):
        raise ValueError(f"Cannot set value at final key {last_key}")

def _update(filename, data, paths):
    basename = os.path.basename(filename)
    basename_no_ext = os.path.splitext(basename)[0]
    new_filename = os.path.join(basename_no_ext, basename)
    os.makedirs(basename_no_ext, exist_ok=True)
    updated = data.copy()
    for path, value in paths:
        if value != "||":
            fname = value.strip("|")
            fullname = _get_fullname(fname)
            if not os.path.isfile(fullname):
                print(f"Warning: file {fullname} does not exist. Skipping.")
                continue
            base_filename = os.path.basename(fullname)
            new_fullname = _existing(basename_no_ext, base_filename)
            
            shutil.copy(fullname, os.path.join(basename_no_ext, new_fullname))
            new_value = f"|{os.path.basename(new_fullname)}|"
            _set_value(updated, path, new_value)
            print(f"Updated path: {' -> '.join(path)}, new value: {new_value}")
    pfs = mikeio.PfsDocument(updated)
    pfs.write(new_filename)

def _zip(folder_path):
    folder_path = os.path.abspath(folder_path)
    parent_dir = os.path.dirname(folder_path)
    folder_name = os.path.basename(folder_path)

    zip_path = shutil.make_archive(
        os.path.join(parent_dir, folder_name),
        "zip",
        root_dir=parent_dir,
        base_dir=folder_name,  # 👈 includes the folder itself
    )

    shutil.rmtree(folder_path)
    return zip_path

def bundle_setup(fname):
    pfs = mikeio.read_pfs(fname)
    results = _find_inputs(pfs.to_dict())
    _update(fname, pfs.to_dict(), results)
    zip_path = _zip(os.path.splitext(fname)[0])
    print(f"Created zip file: {zip_path}")