"""What a session contains, decided by reading the frames rather than the names.

The first version of this pipeline parsed filenames: it expected
`calibrated-T32-...-Luminance-BIN2-W-300-001.fit` and broke on anything else.
That works for exactly one telescope, one binning and one exposure length. Every
other session in the archive - narrowband from 2024, an Ha-only run from 2021,
an RGB set with no luminance at all - fails that pattern.

So this module opens each frame and asks the header. FITS carries FILTER,
EXPTIME, XBINNING, OBJECT and DATE-OBS precisely so that software does not have
to guess from a filename, and iTelescope fills them in properly.

The unit of work is a *session*: one target, one telescope, one night, in one
directory. Everything else in the pipeline takes a Session and asks it what
filters exist and which frames belong to each.
"""
import glob
import json
import os
import zipfile
from collections import defaultdict

from astropy.io import fits

import layout

# Filters that carry broadband colour, in the order a colour image wants them.
BROADBAND = ("Red", "Green", "Blue")
LUMINANCE = ("Luminance", "Clear", "Lum", "L")
# Narrowband, with the aliases different telescopes write into FILTER.
NARROWBAND = {
    "Ha": ("Ha", "H-alpha", "Halpha", "HA", "H_Alpha"),
    "OIII": ("OIII", "O3", "Oiii", "O-III"),
    "SII": ("SII", "S2", "Sii", "S-II"),
}


def canonical_filter(raw):
    """Map whatever the header says onto a name the pipeline knows."""
    if raw is None:
        return "Unknown"
    name = str(raw).strip()
    for canon, aliases in NARROWBAND.items():
        if name in aliases or name.lower() == canon.lower():
            return canon
    for canon in BROADBAND:
        if name.lower() == canon.lower() or name.lower() == canon[0].lower():
            return canon
    if name.lower() in [a.lower() for a in LUMINANCE]:
        return "Luminance"
    return name or "Unknown"


class Frame:
    """One light frame and the header facts the pipeline needs."""

    __slots__ = ("path", "filter", "exptime", "binning", "object", "date_obs",
                 "telescope", "instrument", "shape", "scale", "calibrated")

    def __init__(self, path, header, shape, calibrated=True):
        self.calibrated = calibrated
        self.path = path
        self.filter = canonical_filter(header.get("FILTER"))
        self.exptime = float(header.get("EXPTIME", header.get("EXPOSURE", 0)))
        self.binning = int(header.get("XBINNING", 1))
        self.object = str(header.get("OBJECT", "")).strip()
        self.date_obs = str(header.get("DATE-OBS", ""))
        self.telescope = str(header.get("TELESCOP", "")).strip()
        self.instrument = str(header.get("INSTRUME", "")).strip()
        self.shape = shape
        # iTelescope writes the plate scale in; without it, solve.py estimates.
        scale = header.get("HIERARCH iTelescopePlateScaleH")
        try:
            self.scale = float(scale) if scale is not None else None
        except (TypeError, ValueError):
            self.scale = None

    @property
    def name(self):
        return os.path.basename(self.path)

    def __repr__(self):
        return (f"<{self.filter} {self.exptime:.0f}s bin{self.binning} "
                f"{self.name}>")


class Session:
    """One target, one telescope, one night."""

    def __init__(self, root):
        self.root = os.path.abspath(root)
        self.frames = []
        self._by_filter = {}

    # ---------------------------------------------------------------- ingest
    @staticmethod
    def _is_mirror_junk(path):
        """True for artefacts a broken segmented download leaves behind.

        Cyberduck writes a DIRECTORY named after the file it is fetching, and
        leaves `.cyberducksegment` parts inside when interrupted. Those
        directories are named `something.fit`, so anything that trusts the
        extension will try to open a directory as a frame. A re-mirror over an
        already-organised tree produces dozens of them.
        """
        name = os.path.basename(path).lower()
        if ".cyberducksegment" in name or name.endswith((".part", ".partial")):
            return True
        parent = os.path.basename(os.path.dirname(path)).lower()
        return parent.endswith((".fit", ".fits", ".tif", ".tiff", ".zip"))

    def ingest(self, verbose=True):
        """Put the session into the standard layout and uncompress the lights.

        Handles the three states a downloaded session can be in: still zipped,
        already extracted, or a mixture of both (which is what happens when
        someone unzips a few by hand).
        """
        raw = os.path.join(self.root, layout.RAW)
        cal = os.path.join(self.root, layout.CALIBRATED)
        os.makedirs(raw, exist_ok=True)
        os.makedirs(cal, exist_ok=True)

        moved = extracted = skipped = duplicates = 0
        # Loose files at the session root get filed first.
        for path in sorted(glob.glob(os.path.join(self.root, "*"))):
            if os.path.isdir(path):
                continue
            name = os.path.basename(path)
            low = name.lower()
            if self._is_mirror_junk(path):
                continue
            if low.endswith(".zip") or (low.startswith("jpeg-") and
                                        low.endswith(".jpg")):
                dest = os.path.join(raw, name)
            elif low.endswith((".fit", ".fits", ".tif", ".tiff")):
                dest = os.path.join(cal, name)
            else:
                continue
            if not os.path.exists(dest):
                os.replace(path, dest)
                moved += 1
            elif os.path.getsize(dest) == os.path.getsize(path):
                # Already filed, and the mirror has put an identical copy back.
                os.remove(path)
                duplicates += 1

        # iTelescope ships each frame twice: 'calibrated-*' with bias, dark
        # and flat applied, and 'raw-*' without. Extracting both put two copies
        # of every frame into the stack. Only the calibrated set is extracted;
        # the raw archives stay zipped in raw/ as the untouched originals. If a
        # session has no calibrated archives at all, the raw ones are used and
        # the fact is recorded.
        archives = sorted(glob.glob(os.path.join(raw, "*.zip")))
        calibrated_archives = [z for z in archives
                               if os.path.basename(z).lower().startswith(
                                   "calibrated")]
        self.used_raw = not calibrated_archives and bool(archives)
        if self.used_raw:
            print("    NOTE: no calibrated archives in this session - "
                  "extracting the raw frames instead. They have NOT had bias, "
                  "dark or flat correction applied.")
        else:
            archives = calibrated_archives

        for zpath in archives:
            try:
                with zipfile.ZipFile(zpath) as zf:
                    for info in zf.infolist():
                        base = os.path.basename(info.filename)
                        if not base.lower().endswith((".fit", ".fits", ".tif",
                                                      ".tiff")):
                            continue
                        target = os.path.join(cal, base)
                        if os.path.exists(target) and \
                                os.path.getsize(target) == info.file_size:
                            skipped += 1
                            continue
                        with zf.open(info) as src, open(target, "wb") as dst:
                            while chunk := src.read(1 << 20):
                                dst.write(chunk)
                        extracted += 1
            except zipfile.BadZipFile:
                print(f"    WARNING: {os.path.basename(zpath)} is not a "
                      f"readable zip - skipped")
        if verbose:
            print(f"    filed {moved}, extracted {extracted}, "
                  f"already present {skipped}, "
                  f"duplicate copies removed {duplicates}")
        return dict(moved=moved, extracted=extracted, skipped=skipped,
                    duplicates=duplicates)

    # ----------------------------------------------------------------- scan
    def scan(self, verbose=True):
        """Read every calibrated light frame's header."""
        cal = os.path.join(self.root, layout.CALIBRATED)
        paths = sorted(p for p in glob.glob(os.path.join(cal, "*.fit")) +
                       glob.glob(os.path.join(cal, "*.fits"))
                       if os.path.isfile(p) and not self._is_mirror_junk(p))
        self.frames = []
        for path in paths:
            try:
                with fits.open(path, memmap=True) as hd:
                    hdr = hd[0].header
                    shape = hd[0].shape
            except Exception as exc:                          # noqa: BLE001
                print(f"    WARNING: cannot read {os.path.basename(path)}: "
                      f"{exc}")
                continue
            # Skip anything that is not a light frame.
            if str(hdr.get("IMAGETYP", "Light")).lower().startswith(
                    ("dark", "bias", "flat")):
                continue
            self.frames.append(Frame(path, hdr,
                                     shape, calibrated=bool(
                                         str(hdr.get("CALSTAT", "")).strip())))

        # If both calibrated and uncalibrated frames are present - which
        # happens when a session was extracted before this rule existed - keep
        # only the calibrated ones rather than stacking each frame twice.
        if any(f.calibrated for f in self.frames):
            dropped = [f for f in self.frames if not f.calibrated]
            if dropped and verbose:
                print(f"    ignoring {len(dropped)} uncalibrated duplicate "
                      f"frames (no CALSTAT)")
            self.frames = [f for f in self.frames if f.calibrated]

        self._by_filter = defaultdict(list)
        for f in self.frames:
            self._by_filter[f.filter].append(f)
        if verbose:
            self.describe()
        return self.frames

    # ------------------------------------------------------------ properties
    @property
    def filters(self):
        """Filters present, deepest first - the deepest becomes the reference."""
        return sorted(self._by_filter,
                      key=lambda f: -sum(x.exptime for x in self._by_filter[f]))

    def by_filter(self, name):
        return list(self._by_filter.get(name, []))

    @property
    def target(self):
        names = [f.object for f in self.frames if f.object]
        return max(set(names), key=names.count) if names else \
            os.path.basename(os.path.dirname(self.root))

    @property
    def telescope(self):
        names = [f.telescope for f in self.frames if f.telescope]
        return max(set(names), key=names.count) if names else "unknown"

    @property
    def scale(self):
        vals = [f.scale for f in self.frames if f.scale]
        return sum(vals) / len(vals) if vals else None

    @property
    def palette(self):
        """Which colour scheme this session's filters support."""
        present = set(self._by_filter)
        has_lum = "Luminance" in present
        broad = [c for c in BROADBAND if c in present]
        narrow = [n for n in ("SII", "Ha", "OIII") if n in present]
        if len(broad) == 3 and has_lum:
            return "LRGB"
        if len(broad) == 3:
            return "RGB"
        if len(narrow) >= 3:
            return "SHO"
        if "Ha" in narrow and "OIII" in narrow:
            return "HOO"
        if len(present) == 1:
            return "MONO"
        if broad or narrow:
            return "PARTIAL"
        return "UNKNOWN"

    def describe(self):
        print(f"    target      {self.target}")
        print(f"    telescope   {self.telescope}")
        if self.scale:
            print(f"    plate scale {self.scale:.4f} arcsec/px")
        print(f"    palette     {self.palette}")
        for name in self.filters:
            fs = self._by_filter[name]
            exps = sorted({f.exptime for f in fs})
            bins = sorted({f.binning for f in fs})
            total = sum(f.exptime for f in fs) / 60.0
            shapes = sorted({f.shape for f in fs})
            shape_txt = "/".join(f"{s[1]}x{s[0]}" for s in shapes)
            print(f"      {name:10s} {len(fs):3d} frames  "
                  f"{'/'.join(f'{e:.0f}' for e in exps):>12s} s  "
                  f"bin{'/'.join(str(b) for b in bins)}  "
                  f"{shape_txt:>11s}  {total:6.1f} min")

    # ------------------------------------------------------------- manifest
    def manifest(self):
        return {
            "root": self.root,
            "target": self.target,
            "telescope": self.telescope,
            "plate_scale": self.scale,
            "palette": self.palette,
            "filters": {
                name: {
                    "frames": len(self._by_filter[name]),
                    "total_exposure_s": sum(f.exptime
                                            for f in self._by_filter[name]),
                    "exposures": sorted({f.exptime
                                         for f in self._by_filter[name]}),
                    "binning": sorted({f.binning
                                       for f in self._by_filter[name]}),
                } for name in self.filters
            },
        }

    def save_manifest(self):
        path = os.path.join(self.root, "manifest.json")
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(self.manifest(), fh, indent=2)
        return path


def discover(root):
    """Find every session under a directory tree.

    A session is any directory that directly contains light frames or their
    archives, which matches how iTelescope's downloads are organised
    (telescope / object / date).
    """
    found = []
    for dirpath, dirnames, filenames in os.walk(root):
        base = os.path.basename(dirpath)
        if base.lower().endswith((".fit", ".fits", ".tif", ".tiff", ".zip"))                 or ".cyberducksegment" in base.lower():
            dirnames[:] = []
            continue
        if base in (layout.RAW, layout.CALIBRATED, "intermediates",
                    "final", "renderings", "science", "stacks"):
            dirnames[:] = []
            continue
        has_data = any(f.lower().endswith((".zip", ".fit", ".fits"))
                       for f in filenames)
        has_layout = any(d in dirnames for d in (layout.RAW,
                                                 layout.CALIBRATED))
        if has_data or has_layout:
            found.append(dirpath)
            dirnames[:] = [d for d in dirnames
                           if d in (layout.RAW, layout.CALIBRATED)]
    return sorted(found)