from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Iterable

import numpy as np


def read_json(path: str | Path) -> dict[str, Any]:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def write_json(path: str | Path, payload: Any) -> None:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
    lines = Path(path).read_text(encoding="utf-8").splitlines()
    records: list[dict[str, Any]] = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        records.append(json.loads(line))
    return records


def write_jsonl(path: str | Path, rows: Iterable[dict[str, Any]]) -> None:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    chunks = [json.dumps(row, separators=(",", ":")) for row in rows]
    out_path.write_text("\n".join(chunks) + ("\n" if chunks else ""), encoding="utf-8")


def _to_grayscale(arr: np.ndarray) -> np.ndarray:
    if arr.ndim == 2:
        return arr
    if arr.ndim == 3:
        if arr.shape[2] == 1:
            return arr[..., 0]
        # RGB weighted grayscale conversion.
        rgb = arr[..., :3].astype(np.float32)
        gray = 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
        return gray.astype(arr.dtype if np.issubdtype(arr.dtype, np.integer) else np.float32)
    raise ValueError(f"Expected 2D or 3D image, got shape {arr.shape!r}")


def load_image(path: str | Path) -> np.ndarray:
    image_path = Path(path)
    try:
        from skimage import io as skio  # type: ignore

        arr = skio.imread(str(image_path))
    except Exception:
        # Fallback for environments without scikit-image.
        from matplotlib import image as mpimg  # type: ignore

        arr = mpimg.imread(str(image_path))
    arr = np.asarray(arr)
    arr = _to_grayscale(arr)
    if arr.dtype == np.bool_:
        arr = arr.astype(np.uint8) * 255
    return arr


def save_image(path: str | Path, image: np.ndarray) -> None:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    arr = np.asarray(image)
    try:
        from skimage import io as skio  # type: ignore

        skio.imsave(str(out_path), arr, check_contrast=False)
    except Exception:
        from PIL import Image  # type: ignore

        Image.fromarray(arr).save(str(out_path))


def ensure_time_axis(image_tx: np.ndarray, time_axis: str = "y") -> np.ndarray:
    """Return kymograph in (T, X) orientation."""
    axis = (time_axis or "y").lower()
    if axis not in {"x", "y"}:
        raise ValueError("time_axis must be 'x' or 'y'")
    if image_tx.ndim != 2:
        raise ValueError(f"Expected 2D kymograph, got shape {image_tx.shape!r}")
    return image_tx if axis == "y" else image_tx.T


def percentile_normalize(image: np.ndarray, low: float = 1.0, high: float = 99.0) -> np.ndarray:
    arr = np.asarray(image, dtype=np.float32)
    p_low, p_high = np.percentile(arr, (low, high))
    denom = max(1e-6, float(p_high - p_low))
    out = np.clip((arr - p_low) / denom, 0.0, 1.0)
    return out.astype(np.float32)


def resolve_manifest_path(base_dir: str | Path, value: str | Path) -> Path:
    p = Path(value)
    return p if p.is_absolute() else Path(base_dir) / p
