"""
reader.py — core metadata extraction from Zeiss CZI files.
"""

from __future__ import annotations

import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

# ── optional CZI readers ──────────────────────────────────────────────────────
try:
    from aicspylibczi import CziFile as _AicsCziFile
    _HAS_AICS = True
except ImportError:
    _HAS_AICS = False

try:
    import czifile as _czifile
    _HAS_CZIFILE = True
except ImportError:
    _HAS_CZIFILE = False


# ── public data class ─────────────────────────────────────────────────────────

@dataclass
class CziMetadata:
    """All metadata extracted from a CZI file by czipeak."""

    # file
    path:    str   = ""
    size_mb: float = 0.0

    # acquisition
    acq_date_raw: Optional[str] = None
    objective:    Optional[str] = None
    magnification: Optional[str] = None

    # binary dims
    axes:  str   = ""
    shape: tuple = ()

    # temporal
    t_count:               Optional[int]      = None
    t_start:               Optional[datetime] = None
    t_stop:                Optional[datetime] = None
    total_duration_s:      Optional[float]    = None
    dt_actual_s:           Optional[float]    = None
    interval_configured_s:   Optional[float]  = None
    interval_configured_raw: Optional[str]    = None

    # tiles / scenes
    s_count: Optional[int]  = None
    scenes:  List[Dict]     = field(default_factory=list)

    # Z-stack
    z_count:    Optional[int]   = None
    dz_um:      Optional[float] = None
    z_start_um: Optional[float] = None
    z_end_um:   Optional[float] = None
    z_range_um: Optional[float] = None

    # channels
    channels: List[Dict] = field(default_factory=list)

    # XY pixel size
    dx_um: Optional[float] = None
    dy_um: Optional[float] = None


# ── XML helpers ───────────────────────────────────────────────────────────────

def _find_text(root: ET.Element, *tags: str) -> Optional[str]:
    for tag in tags:
        el = root.find(tag)
        if el is not None and el.text is not None:
            return el.text.strip()
    return None


def _parse_iso(s: str) -> Optional[datetime]:
    s = s.strip()
    if "." in s and s.endswith("Z"):
        base, frac = s[:-1].split(".", 1)
        s = base + "." + frac[:6] + "Z"
    try:
        return datetime.strptime(s, "%Y-%m-%dT%H:%M:%S.%fZ").replace(
            tzinfo=timezone.utc)
    except ValueError:
        return None


def _timespan_to_seconds(value: float, unit: str) -> float:
    u = unit.strip().lower()
    if u in ("us", "µs", "microsecond", "microseconds"):  return value / 1e6
    if u in ("ms", "millisecond", "milliseconds"):         return value / 1e3
    if u in ("s", "sec", "second", "seconds"):             return value
    if u in ("min", "minute", "minutes"):                  return value * 60.0
    if u in ("h", "hr", "hour", "hours"):                  return value * 3600.0
    return value  # fallback: assume seconds


# ── XML loader ────────────────────────────────────────────────────────────────

def _load_xml(czi_path: str) -> ET.Element:
    if _HAS_AICS:
        try:
            czi = _AicsCziFile(czi_path)
            return ET.fromstring(ET.tostring(czi.meta, encoding="unicode"))
        except Exception:
            pass
    if _HAS_CZIFILE:
        with _czifile.CziFile(czi_path) as czi:
            raw = czi.metadata()
            xml_str = raw.decode("utf-8") if isinstance(raw, bytes) else raw
            return ET.fromstring(xml_str)
    raise ImportError(
        "No CZI reader found.\n"
        "Install one of:\n"
        "  pip install aicspylibczi lxml\n"
        "  pip install czifile"
    )


# ── per-section extractors ────────────────────────────────────────────────────

def _extract_binary_dims(czi_path: str) -> Dict:
    out: Dict = {}
    if not _HAS_CZIFILE:
        return out
    try:
        with _czifile.CziFile(czi_path) as czi:
            axes  = czi.axes   # plain str e.g. 'STCYX0'
            shape = czi.shape
            out["axes"]  = axes
            out["shape"] = tuple(shape)
            for dim in ("T", "S", "Z", "C"):
                if dim in axes:
                    out[f"{dim}_count"] = shape[axes.index(dim)]
    except Exception:
        pass
    return out


def _extract_timestamp_list(root: ET.Element) -> List[float]:
    """Return per-frame timestamps (s, relative to t=0) from legacy schema."""
    ts: List[float] = []

    block = root.find(".//Dimensions/T/TimeStamps")
    if block is not None:
        for el in block.findall("TimeStamp"):
            txt = el.text or el.get("Value", "")
            if txt:
                try: ts.append(float(txt))
                except ValueError: pass

    if not ts:
        for el in root.iter("TimeStamp"):
            txt = el.text or el.get("Value", "")
            if txt:
                try: ts.append(float(txt))
                except ValueError: pass

    if not ts:
        for el in root.iter("RelativeTime"):
            txt = el.text or ""
            if txt:
                try: ts.append(float(txt))
                except ValueError: pass

    if not ts:
        return ts
    if max(ts) > 1e7:           # Windows 100-ns ticks → seconds
        ts = [t / 1e7 for t in ts]
    t0 = ts[0]
    return [t - t0 for t in ts]


def _extract_temporal(root: ET.Element, binary: Dict) -> Dict:
    out: Dict = {}
    if "T_count" in binary:
        out["t_count"] = binary["T_count"]

    start_el = root.find(".//Dimensions/T/StartTime")
    if start_el is not None and start_el.text is not None:
        ts = _parse_iso(start_el.text)
        if ts:
            out["t_start"] = ts

    for el in root.iter("TimelineElement"):
        if (el.findtext(".//Type") or "").strip().lower() == "stop":
            time_el = el.find("Time")
            if time_el is not None and time_el.text is not None:
                ts = _parse_iso(time_el.text)
                if ts:
                    out["t_stop"] = ts
            bounds = el.find("Bounds")
            if bounds is not None and bounds.get("StartT") is not None:
                out["last_t_index"] = int(str(bounds.get("StartT")))
            break

    if "t_start" in out and "t_stop" in out:
        total_s = (out["t_stop"] - out["t_start"]).total_seconds()
        out["total_duration_s"] = total_s
        n = out.get("t_count") or (out.get("last_t_index", 0) + 1)
        if n and n > 1:
            out["dt_actual_s"] = total_s / (n - 1)

    ts_setup = root.find(".//TimeSeriesSetup")
    if ts_setup is not None:
        iv_val_el  = ts_setup.find(".//Interval/TimeSpan/Value")
        iv_unit_el = ts_setup.find(".//Interval/TimeSpan/DefaultUnitFormat")
        if iv_val_el is not None and iv_val_el.text is not None:
            iv_val  = float(iv_val_el.text)
            iv_unit = str(iv_unit_el.text) if (
                iv_unit_el is not None and iv_unit_el.text) else "s"
            out["interval_configured_s"]   = _timespan_to_seconds(iv_val, iv_unit)
            out["interval_configured_raw"] = f"{iv_val} {iv_unit}"

    return out


def _extract_z_info(root: ET.Element, binary: Dict) -> Dict:
    """
    Extract Z-stack info. Returns empty dict if no Z axis.

    Keys: z_count, dz_um, z_start_um, z_end_um, z_range_um
    """
    if binary.get("Z_count") is None:
        return {}
    out: Dict = {"z_count": binary["Z_count"]}

    # dz from <Scaling/Items/Distance Id="Z"> in metres
    for dist in root.findall(".//Scaling/Items/Distance"):
        if dist.get("Id") == "Z":
            val_el = dist.find("Value")
            if val_el is not None and val_el.text is not None:
                out["dz_um"] = float(val_el.text) * 1e6
            break

    # Z start/increment from <Dimensions/Z/Positions/Interval> in µm
    z_dim = root.find(".//Dimensions/Z")
    if z_dim is not None:
        start_el = z_dim.find(".//Positions/Interval/Start")
        incr_el  = z_dim.find(".//Positions/Interval/Increment")
        if start_el is not None and start_el.text is not None:
            out["z_start_um"] = float(start_el.text)
        if incr_el is not None and incr_el.text is not None and "dz_um" not in out:
            out["dz_um"] = float(incr_el.text)

    n     = out.get("z_count", 0)
    dz    = out.get("dz_um")
    start = out.get("z_start_um")
    if n and dz and start is not None:
        out["z_end_um"]   = start + dz * (n - 1)
        out["z_range_um"] = dz * (n - 1)

    return out


def _extract_scenes(root: ET.Element, binary: Dict) -> Dict:
    out: Dict = {}
    if "S_count" in binary:
        out["s_count"] = binary["S_count"]
    scene_els = root.findall(".//Dimensions/S/Scenes/Scene")
    if scene_els:
        out["scenes"] = [
            {
                "index":  s.get("Index", str(i)),
                "name":   s.get("Name",  f"Scene{i}"),
                "center": (s.findtext("CenterPosition") or "").strip(),
            }
            for i, s in enumerate(scene_els)
        ]
    return out


def _extract_channels(root: ET.Element) -> List[Dict]:
    return [
        {
            "index": i,
            "name":  ch.get("Name") or ch.findtext("Name") or f"Ch{i}",
            "fluor": ch.findtext("Fluor") or ch.findtext(".//FluorescenceDye") or "",
            "ex_nm": ch.findtext("ExcitationWavelength") or
                     ch.findtext(".//ExcitationWavelength") or "",
            "em_nm": ch.findtext("EmissionWavelength") or
                     ch.findtext(".//EmissionWavelength") or "",
        }
        for i, ch in enumerate(root.findall(".//Channels/Channel"))
    ]


def _extract_pixel_size(root: ET.Element) -> Dict:
    out: Dict = {}
    for dist in root.findall(".//Scaling/Items/Distance"):
        dim_id = dist.get("Id", "")
        val_el = dist.find("Value")
        if val_el is not None and val_el.text is not None:
            val_m = float(val_el.text)
            if   dim_id == "X": out["dx_um"] = val_m * 1e6
            elif dim_id == "Y": out["dy_um"] = val_m * 1e6
    return out


# ── main public function ──────────────────────────────────────────────────────

def read_czi(czi_path: str) -> CziMetadata:
    """
    Read metadata from a Zeiss CZI file.

    Parameters
    ----------
    czi_path : str or Path
        Path to the .czi file.

    Returns
    -------
    CziMetadata
        All extracted fields. Fields absent in the file are ``None`` or empty.
    """
    path = Path(czi_path)
    meta = CziMetadata(path=str(path), size_mb=path.stat().st_size / 1e6)

    root   = _load_xml(str(path))
    binary = _extract_binary_dims(str(path))

    meta.axes  = binary.get("axes", "")
    meta.shape = binary.get("shape", ())

    meta.acq_date_raw  = _find_text(root,
        ".//Experiment/AcquisitionDateAndTime",
        ".//AcquisitionDateAndTime", ".//StartTime")
    meta.objective     = _find_text(root,
        ".//Objectives/Objective/Name", ".//Objective/Name")
    meta.magnification = _find_text(root,
        ".//Objectives/Objective/NominalMagnification", ".//NominalMagnification")

    temp = _extract_temporal(root, binary)
    meta.t_count               = temp.get("t_count")
    meta.t_start               = temp.get("t_start")
    meta.t_stop                = temp.get("t_stop")
    meta.total_duration_s      = temp.get("total_duration_s")
    meta.dt_actual_s           = temp.get("dt_actual_s")
    meta.interval_configured_s = temp.get("interval_configured_s")
    meta.interval_configured_raw = temp.get("interval_configured_raw")

    sc = _extract_scenes(root, binary)
    meta.s_count = sc.get("s_count")
    meta.scenes  = sc.get("scenes", [])

    z = _extract_z_info(root, binary)
    meta.z_count    = z.get("z_count")
    meta.dz_um      = z.get("dz_um")
    meta.z_start_um = z.get("z_start_um")
    meta.z_end_um   = z.get("z_end_um")
    meta.z_range_um = z.get("z_range_um")

    meta.channels = _extract_channels(root)

    px = _extract_pixel_size(root)
    meta.dx_um = px.get("dx_um")
    meta.dy_um = px.get("dy_um")

    return meta
