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

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):
    """Return a GeoDataFrame of every element with x/y indices, in the most
    appropriate CRS for the file (see module docstring)."""
    g = mikeio.open(dfs2_path).geometry
    nx, ny = g.nx, g.ny
    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(), "geometry": geom},
        crs=crs,
    )
 
 
def dfs2_to_shapefile(input_fname, output_fname, verbose=True):
    """Write every element of `input_fname` as a polygon to `output_fname`."""
    gdf = dfs2_to_geodataframe(input_fname)
    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 to {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