import mikeio
import os
import geopandas as gpd
from shapely import polygons
import numpy as np
import re
import tempfile
import zipfile


from ._utils import _element_corners, _read_text_preserving, PackageReport, _write_text_preserving, _resolve, _identity, _crosscheck_with_mikeio
from ._utils import _PIPE, _resolve_item, _resolve_time_index

def mesh2poly(input_fname, output_fname, verbose=True):
    """
    Convert a MIKE mesh file to a GeoDataFrame of polygons.
    Parameters:
    - input_fname: Path to the MIKE mesh file (.mesh).
    - output_fname: Path to save the output GeoDataFrame (.shp).
    - verbose: If True, print progress messages.
    Returns:
    - A GeoDataFrame containing the polygons of the mesh elements.
    """
    # Check if file exists
    if not os.path.exists(input_fname):
        raise FileNotFoundError(f"Input file not found: {input_fname}")
    # Check file extensions
    if not input_fname.endswith(".mesh"):
        raise ValueError("Input file must have a .mesh extension")
    if not output_fname.endswith(".shp"):
        raise ValueError("Output file must have a .shp extension")
    # Make sure output directory exists
    output_dir = os.path.dirname(output_fname)
    if output_dir and not os.path.exists(output_dir):
        if verbose:
            print(f"WARNING: Output directory does not exist. Creating: {output_dir}")
        os.makedirs(output_dir)
    # Load mesh
    mesh = mikeio.Mesh(input_fname)
    geom = mesh.geometry
    projection = geom.projection if geom.projection != "LONG/LAT" else "EPSG:4326"
    z = mesh.element_coordinates[:, 2]
    mpoly = geom.to_shapely()
    polygons = list(mpoly.geoms)
    values = z.flatten()

    gdf = gpd.GeoDataFrame(data={"ID": range(len(polygons)), "Bathymetry": values}, geometry=polygons, crs=projection)
    gdf.to_file(output_fname)
    if verbose:
        print(f"Mesh converted to polygons and saved to: {output_fname}")
 
 
def dfs2_to_geodataframe(dfs2_path, item=0, timestep=0):
    """Return a GeoDataFrame of every element with x/y indices and the value of
    one item at one timestep, in the most appropriate CRS for the file (see
    module docstring).

    Parameters
    ----------
    dfs2_path : str
        Path to the dfs2 file.
    item : int or str
        Item to export, given as an integer index or an item name. Raises
        ValueError (listing the available items) if it does not resolve.
    timestep : int or str/datetime
        Either an integer time index, or an actual timestamp (e.g.
        "2018-01-01 01:00"). A timestamp outside the file range, or one that
        falls between available steps, raises ValueError describing the range,
        the timestep interval, and the nearest available steps.

    The item's values are written to a column named after the item; the
    element order matches x_index/y_index (a mikeio (ny, nx) array raveled
    with x varying fastest).
    """
    dfs = mikeio.open(dfs2_path)
    g = dfs.geometry
    nx, ny = g.nx, g.ny

    # Resolve item (by index or name) and timestep (by index or timestamp).
    iidx, item_name = _resolve_item(dfs.items, item)
    tidx = _resolve_time_index(dfs.time, timestep)
    da = mikeio.read(dfs2_path, items=[iidx], time=tidx)[0]
    # reshape (not squeeze) so a 1-cell-wide grid keeps both axes.
    values = np.asarray(da.to_numpy(), dtype=float).reshape(ny, nx)

    X, Y, crs = _element_corners(g)

    # Corner arrays for every cell, each shape (ny, nx). Clockwise ring order:
    # bottom-left, top-left, top-right, bottom-right (+ closing point).
    def corner(js, is_):
        return np.stack([X[js, is_], Y[js, is_]], axis=-1)  # (ny, nx, 2)

    bl = corner(slice(0, ny),     slice(0, nx))
    tl = corner(slice(1, ny + 1), slice(0, nx))
    tr = corner(slice(1, ny + 1), slice(1, nx + 1))
    br = corner(slice(0, ny),     slice(1, nx + 1))

    rings = np.stack([bl, tl, tr, br, bl], axis=2).reshape(nx * ny, 5, 2)
    geom = polygons(rings)

    jj, ii = np.meshgrid(np.arange(ny), np.arange(nx), indexing="ij")
    return gpd.GeoDataFrame(
        {
            "x_index": ii.ravel(),
            "y_index": jj.ravel(),
            item_name: values.ravel(),
            "geometry": geom,
        },
        crs=crs,
    )


def dfs2_to_shapefile(input_fname, output_fname, item=0, timestep=0, verbose=True):
    """Write every element of `input_fname` as a polygon to `output_fname`,
    carrying x_index, y_index, and the value of `item` at `timestep`.

    `item` accepts an index or item name; `timestep` accepts an index or a
    timestamp string (see `dfs2_to_geodataframe`).
    """
    gdf = dfs2_to_geodataframe(input_fname, item=item, timestep=timestep)

    # Shapefile (.dbf) field names are capped at 10 characters. Rename the
    # value column up-front so we control the truncation instead of letting
    # fiona do it silently (and possibly clash with x_index/y_index).
    value_col = next(c for c in gdf.columns
                     if c not in ("x_index", "y_index", "geometry"))
    if len(value_col) > 10:
        safe = value_col[:10]
        gdf = gdf.rename(columns={value_col: safe})
        if verbose:
            print(f"Item name '{value_col}' truncated to '{safe}' "
                  "(shapefile 10-character field-name limit).")
        value_col = safe

    gdf.to_file(output_fname)  # geopandas writes the .prj from the GeoDataFrame CRS
    if verbose:
        epsg = gdf.crs.to_epsg()
        crs_desc = f"EPSG:{epsg}" if epsg else gdf.crs.name
        print(f"Wrote {len(gdf)} elements (item '{value_col}') to "
              f"{output_fname}  (CRS: {crs_desc})")


# --------------------------------------------------------------------------- #
# Link discovery (usable on its own)
# --------------------------------------------------------------------------- #
def find_linked_files(setup_file: str) -> list[tuple[str, str, bool]]:
    """Return every linked input in ``setup_file``.
 
    Each entry is ``(stored_path, resolved_path, exists)``. Duplicates are kept
    (one entry per occurrence) so you can see how often each file is referenced.
    """
    setup_file = os.path.abspath(setup_file)
    setup_dir = os.path.dirname(setup_file)
    text, _, _ = _read_text_preserving(setup_file)
    out: list[tuple[str, str, bool]] = []
    for m in _PIPE.finditer(text):
        stored = m.group(1)
        if not stored:  # empty ``||`` -> no file linked
            continue
        resolved = _resolve(stored, setup_dir)
        out.append((stored, resolved, os.path.isfile(resolved)))
    return out


def package_mike_setup(
    setup_file: str,
    output_zip: str | None = None,
    *,
    on_missing: str = "error",
    folder_name: str | None = None,
    verify_with_mikeio: bool = False,
) -> PackageReport:
    """Package a MIKE setup and all of its linked inputs into a zip.
 
    Parameters
    ----------
    setup_file
        Path to the ``.m21fm`` / ``.m3fm`` / other PFS setup file.
    output_zip
        Destination zip path. Defaults to ``<setup_dir>/<setup_stem>.zip``.
    on_missing : {"error", "warn"}
        ``"error"`` (default) aborts if any linked file is absent — nothing is
        written. ``"warn"`` proceeds: absent files are not copied and their link
        is left unchanged (correct for inert defaults like ``closuredepth.dfs1``).
    folder_name
        Name of the folder inside the zip that holds the setup + inputs.
        Defaults to the setup file's stem.
    verify_with_mikeio
        If True and mikeio is installed, cross-check that the links we found
        match the file paths mikeio parses. Purely diagnostic; never used to
        write the file.
 
    Returns
    -------
    PackageReport
 
    Notes
    -----
    The original ``setup_file`` is only read, never modified. The zip contains::
 
        <folder_name>/<original setup filename>   (with links rewritten)
        <folder_name>/<input 1>
        <folder_name>/<input 2>
        ...
 
    Rewritten links are bare filenames, which MIKE resolves relative to the
    setup file — so the package runs from wherever it is unzipped.
    """
    if on_missing not in ("error", "warn"):
        raise ValueError("on_missing must be 'error' or 'warn'")
 
    setup_file = os.path.abspath(setup_file)
    if not os.path.isfile(setup_file):
        raise FileNotFoundError(f"Setup file not found: {setup_file}")
 
    setup_dir = os.path.dirname(setup_file)
    setup_basename = os.path.basename(setup_file)
    stem = os.path.splitext(setup_basename)[0]
    folder = folder_name or stem
    if output_zip is None:
        output_zip = os.path.join(setup_dir, stem + ".zip")
    output_zip = os.path.abspath(output_zip)
 
    text, encoding, had_bom = _read_text_preserving(setup_file)
    report = PackageReport(setup_name=setup_basename, folder_name=folder)
 
    # -- Pass 1: discover links, classify existence, plan destinations --------
    identity_to_dest: dict[str, str] = {}   # fs identity -> destination filename
    identity_to_src: dict[str, str] = {}    # fs identity -> a resolved source path
    used_dest_names: set[str] = set()
    ref_counts: dict[str, int] = {}         # fs identity -> occurrences
 
    missing: dict[str, str] = {}
 
    for m in _PIPE.finditer(text):
        stored = m.group(1)
        if not stored:
            continue
        resolved = _resolve(stored, setup_dir)
        if not os.path.isfile(resolved):
            missing[stored] = resolved
            continue
        ident = _identity(resolved)
        ref_counts[ident] = ref_counts.get(ident, 0) + 1
        if ident in identity_to_dest:
            continue  # destination already assigned for this physical file
        # Assign a collision-free destination name (flat folder).
        base = os.path.basename(resolved)
        dest = base
        if dest in used_dest_names:
            root, ext = os.path.splitext(base)
            n = 1
            while f"{root}_{n}{ext}" in used_dest_names:
                n += 1
            dest = f"{root}_{n}{ext}"
            report.renamed_for_collision[dest] = resolved
        used_dest_names.add(dest)
        identity_to_dest[ident] = dest
        identity_to_src[ident] = resolved
 
    # Abort early on missing links if requested — write nothing.
    if missing and on_missing == "error":
        listing = "\n".join(f"  - {s}  (looked in: {r})" for s, r in missing.items())
        raise FileNotFoundError(
            f"{len(missing)} linked file(s) referenced by '{setup_basename}' "
            f"were not found:\n{listing}\n"
            "Fix the paths, or call with on_missing='warn' to package anyway "
            "(missing links are then left unchanged — appropriate for inert "
            "MIKE defaults such as 'closuredepth.dfs1')."
        )
 
    report.missing = missing
    report.copied = {identity_to_src[i]: d for i, d in identity_to_dest.items()}
    report.duplicate_links_collapsed = sum(1 for c in ref_counts.values() if c > 1)
 
    # Optional diagnostic cross-check against mikeio's own parse.
    if verify_with_mikeio:
        _crosscheck_with_mikeio(setup_file, report)
 
    # -- Pass 2: rewrite only the pipe tokens, preserving all other bytes -----
    def _sub(match: re.Match) -> str:
        stored = match.group(1)
        if not stored:
            return match.group(0)  # ``||`` untouched
        resolved = _resolve(stored, setup_dir)
        if not os.path.isfile(resolved):
            return match.group(0)  # missing & tolerated -> leave as-is
        dest = identity_to_dest[_identity(resolved)]
        report.links_rewritten += 1
        return f"|{dest}|"
 
    new_text = _PIPE.sub(_sub, text)
 
    # -- Build the zip via a temp staging copy of the setup (original safe) ---
    os.makedirs(os.path.dirname(output_zip) or ".", exist_ok=True)
    with tempfile.TemporaryDirectory() as tmp:
        staged_setup = os.path.join(tmp, setup_basename)
        _write_text_preserving(staged_setup, new_text, encoding, had_bom)
 
        with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as zf:
            zf.write(staged_setup, arcname=f"{folder}/{setup_basename}")
            for ident, dest in identity_to_dest.items():
                zf.write(identity_to_src[ident], arcname=f"{folder}/{dest}")
 
    report.zip_path = output_zip
    return report