"""Resolve and optionally download one TESS product set from MAST.

Given a TIC identifier and a sector, this script queries the read-only MAST CAOM
services for TESS timeseries and target-pixel products, selects at most one
product per requested type, and prints or writes the resolved URLs. With
``--download`` it fetches only HTTPS URLs on the MAST/STScI allowlist, enforces a
per-file byte cap, hashes each file, and appends a record to the provenance
manifest.

The ``--self-test`` path is offline. It reads a committed synthetic MAST
response fixture, uses a stub opener, writes a temporary manifest with two
products, and never opens a network connection.

Run from the repository root:

    python research/tess-timing/scripts/fetch_products.py --tic-id TIC_ID --sector SECTOR
    python research/tess-timing/scripts/fetch_products.py --tic-id TIC_ID --sector SECTOR --download
    python research/tess-timing/scripts/fetch_products.py --self-test
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import tempfile
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Self

PACKAGE_ROOT = Path(__file__).resolve().parents[1]
REPO_ROOT = PACKAGE_ROOT.parents[1]
SRC = PACKAGE_ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from tess_timing.model import utc_now_iso
from tess_timing.product_fetch import (
    ALLOWED_HOSTS,
    DEFAULT_MAX_BYTES,
    MAST_INVOKE_URL,
    PRODUCT_TYPES,
    ByteCapExceededError,
    DigestMismatchError,
    DuplicateProductError,
    HostNotAllowedError,
    InsecureUrlError,
    NoMatchingProductError,
    Product,
    ProductError,
    download_product,
    require_cadence,
    resolve_products,
)
from tess_timing.provenance import (
    DEFAULT_CODE_VERSION,
    ProvenanceError,
    append_record,
    record_from_file,
    verify_manifest,
)

TASK = "jdp-tess-fetch-products"
FIXTURE_PATH = (
    PACKAGE_ROOT / "data" / "fetch-products-run" / "mast_products_fixture.json"
)
DEFAULT_OUT = PACKAGE_ROOT / "data" / "fetch-products-run" / "run.json"
DEFAULT_MANIFEST = PACKAGE_ROOT / "data" / "fetch-products-run" / "manifest.jsonl"
DEFAULT_DOWNLOADS = PACKAGE_ROOT / "data" / "downloads"
USER_AGENT = "jdp-tess-products/0.1 (repository research script)"
MAST_COLUMNS = "obsid,target_name,sequence_number,dataproduct_type,obs_collection"
QUERY_PAGE_SIZE = 1000


class MastQueryError(ProductError):
    """A MAST CAOM metadata query failed."""


@dataclass(frozen=True)
class MastCall:
    service: str
    request: dict[str, Any]
    requested_utc: str
    http_status: int
    rows: list[dict[str, Any]]

    def record(self) -> dict[str, Any]:
        return {
            "service": self.service,
            "endpoint": MAST_INVOKE_URL,
            "request": self.request,
            "requested_utc": self.requested_utc,
            "http_status": self.http_status,
            "rows": len(self.rows),
        }


def _display(path: Path) -> str:
    try:
        return path.resolve().relative_to(REPO_ROOT.resolve()).as_posix()
    except ValueError:
        return str(path)


def _post_mast(
    service: str,
    params: dict[str, Any],
    *,
    opener: Any | None = None,
    timeout: int = 120,
) -> MastCall:
    request_payload = {
        "service": service,
        "params": params,
        "format": "json",
        "pagesize": QUERY_PAGE_SIZE,
        "page": 1,
    }
    data = urllib.parse.urlencode({"request": json.dumps(request_payload)}).encode(
        "utf-8"
    )
    request = urllib.request.Request(
        MAST_INVOKE_URL,
        data=data,
        headers={
            "User-Agent": USER_AGENT,
            "Content-Type": "application/x-www-form-urlencoded",
            "Accept": "application/json",
        },
    )
    requested_utc = utc_now_iso()
    open_callable = urllib.request.urlopen if opener is None else opener
    try:
        with open_callable(request, timeout=timeout) as response:
            status = int(getattr(response, "status", 200) or 200)
            body = response.read()
    except urllib.error.HTTPError as exc:
        raise MastQueryError(f"{service}: HTTP {exc.code} {exc.reason}") from exc
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        raise MastQueryError(f"{service}: network error: {exc}") from exc

    text = body.decode("utf-8", errors="replace")
    try:
        payload = json.loads(text)
    except json.JSONDecodeError as exc:
        raise MastQueryError(f"{service}: response was not JSON: {exc.msg}") from exc
    if not isinstance(payload, dict):
        raise MastQueryError(f"{service}: response was not a JSON object")
    rows = payload.get("data", [])
    if not isinstance(rows, list):
        raise MastQueryError(f"{service}: response data was not a list")
    return MastCall(
        service=service,
        request=request_payload,
        requested_utc=requested_utc,
        http_status=status,
        rows=rows,
    )


def query_mast_products(
    tic_id: int,
    sector: int,
    *,
    opener: Any | None = None,
    timeout: int = 120,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Query MAST for the TESS observations and products of one TIC sector."""

    observation_call = _post_mast(
        "Mast.Caom.Filtered",
        {
            "columns": MAST_COLUMNS,
            "filters": [
                {"paramName": "obs_collection", "values": ["TESS"]},
                {"paramName": "dataproduct_type", "values": ["timeseries"]},
                {"paramName": "target_name", "values": [str(tic_id)]},
                {"paramName": "sequence_number", "values": [sector]},
            ],
        },
        opener=opener,
        timeout=timeout,
    )
    observations: dict[int, dict[str, Any]] = {}
    for row in observation_call.rows:
        if not isinstance(row, dict):
            continue
        value = row.get("obsid")
        if value in (None, ""):
            continue
        try:
            observations.setdefault(int(value), row)
        except (TypeError, ValueError):
            continue
    unique_obsids = sorted(observations)

    calls = [observation_call]
    product_rows: list[dict[str, Any]] = []
    for obsid in unique_obsids:
        call = _post_mast(
            "Mast.Caom.Products",
            {"obsid": str(obsid)},
            opener=opener,
            timeout=timeout,
        )
        calls.append(call)
        observation = observations[obsid]
        for row in call.rows:
            if not isinstance(row, dict):
                continue
            product = dict(row)
            for field in ("target_name", "sequence_number", "dataproduct_type"):
                if product.get(field) in (None, ""):
                    product[field] = observation.get(field)
            product_rows.append(product)

    record = {
        "service": "MAST CAOM",
        "endpoint": MAST_INVOKE_URL,
        "requested_utc": calls[0].requested_utc,
        "observation_rows": len(observation_call.rows),
        "product_rows": len(product_rows),
        "calls": [call.record() for call in calls],
    }
    return product_rows, record


def _fixture_rows(path: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ProductError(f"cannot read fixture {path}: {exc}") from exc
    if not isinstance(data, dict):
        raise ProductError("fixture must be a JSON object")
    response = data.get("response")
    if not isinstance(response, dict) or not isinstance(response.get("data"), list):
        raise ProductError("fixture has no response.data list")
    rows = [row for row in response["data"] if isinstance(row, dict)]
    metadata = data.get("_fixture") if isinstance(data.get("_fixture"), dict) else {}
    record = {
        "service": "fixture",
        "fixture": _display(path),
        "requested_utc": utc_now_iso(),
        "observation_rows": 0,
        "product_rows": len(rows),
        "calls": [],
        "fixture_synthetic": bool(metadata.get("synthetic", False)),
    }
    return rows, record


def _product_record(product: Product) -> dict[str, Any]:
    return {
        "product_type": product.product_type,
        "product_subgroup": product.product_subgroup,
        "obsid": product.obsid,
        "obs_id": product.obs_id,
        "target_name": product.target_name,
        "sector": product.sector,
        "cadence": product.cadence,
        "product_filename": product.filename,
        "file_url": product.file_url,
        "declared_bytes": product.declared_bytes,
    }


def _write_json(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    text = json.dumps(data, indent=2, ensure_ascii=False) + "\n"
    path.write_text(text, encoding="utf-8")


class _StubResponse:
    def __init__(
        self,
        payload: bytes,
        url: str,
        *,
        status: int = 200,
        include_length: bool = True,
    ) -> None:
        self._payload = payload
        self._url = url
        self._offset = 0
        self.status = status
        self.headers = {"Content-Length": str(len(payload))} if include_length else {}

    def read(self, size: int = -1) -> bytes:
        if size is None or size < 0:
            end = len(self._payload)
        else:
            end = min(len(self._payload), self._offset + size)
        chunk = self._payload[self._offset : end]
        self._offset = end
        return chunk

    def geturl(self) -> str:
        return self._url

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type, exc, traceback) -> bool:
        return False


class _StubOpener:
    def __init__(self, payloads: dict[str, bytes]) -> None:
        self._payloads = payloads

    def open(
        self, request: urllib.request.Request, timeout: int | None = None
    ) -> _StubResponse:
        url = request.full_url
        if url not in self._payloads:
            raise AssertionError(f"unexpected stub URL: {url}")
        return _StubResponse(self._payloads[url], url)


def _check(condition: bool, message: str) -> None:
    if not condition:
        raise ProductError(f"self-test failed: {message}")


def _expect_error(label: str, error_type: type[ProductError], action) -> None:
    try:
        action()
    except error_type:
        return
    except ProductError as exc:
        raise ProductError(
            f"self-test failed: {label}: expected {error_type.__name__}, "
            f"got {type(exc).__name__}: {exc}"
        ) from exc
    raise ProductError(
        f"self-test failed: {label}: expected {error_type.__name__}, none was raised"
    )


def _run_self_test() -> int:
    rows, query_record = _fixture_rows(FIXTURE_PATH)
    fixture = json.loads(FIXTURE_PATH.read_text(encoding="utf-8"))
    metadata = fixture["_fixture"]
    tic_id = int(metadata["tic_id"])
    sector = int(metadata["sector"])
    _check(metadata.get("synthetic") is True, "fixture is not labelled synthetic")
    _check(query_record["fixture_synthetic"] is True, "fixture metadata was not read")

    max_bytes = 4096
    products = resolve_products(rows, tic_id=tic_id, sector=sector, max_bytes=max_bytes)
    _check(len(products) == 2, f"expected two products, found {len(products)}")
    _check(
        {product.product_type for product in products}
        == {"lightcurve", "target-pixel-file"},
        "the two products are not one light curve and one target pixel file",
    )

    payloads = {
        "lightcurve": b"synthetic TESS light curve fixture\n" * 20,
        "target-pixel-file": b"synthetic TESS target pixel fixture\n" * 24,
    }
    stub = _StubOpener(
        {product.file_url: payloads[product.product_type] for product in products}
    )

    scratch_root = REPO_ROOT / ".scratch"
    scratch_root.mkdir(parents=True, exist_ok=True)
    checks = 0
    with tempfile.TemporaryDirectory(
        prefix="fetch-products-self-test-", dir=scratch_root
    ) as temporary_directory:
        work = Path(temporary_directory)
        manifest = work / "manifest.jsonl"
        manifest_urls: list[str] = []
        product_entries: list[dict[str, Any]] = []
        for product in products:
            destination = (
                work / "downloads" / str(tic_id) / str(sector) / product.filename
            )
            result = download_product(
                product,
                destination,
                max_bytes=max_bytes,
                opener=stub.open,
            )
            record = record_from_file(
                destination,
                target=product.target_name,
                tic_id=tic_id,
                sector=sector,
                cadence=require_cadence(product),
                product_type=product.product_type,
                file_url=product.file_url,
                aperture="not applied",
                quality_mask="not applied",
                code_version=DEFAULT_CODE_VERSION,
                downloaded_utc=result.downloaded_utc,
            )
            append_record(manifest, record, file_path=destination)
            manifest_urls.append(product.file_url)
            product_entries.append(
                {
                    **_product_record(product),
                    "http_status": result.http_status,
                    "requested_utc": result.requested_utc,
                    "downloaded_utc": result.downloaded_utc,
                    "byte_count": result.byte_count,
                    "sha256": result.sha256,
                    "local_path": _display(destination),
                }
            )
            checks += 1

        records = verify_manifest(manifest)
        _check(len(records) == 2, "the manifest does not hold two records")
        for record in records:
            expected_digest = hashlib.sha256(payloads[record.product_type]).hexdigest()
            _check(
                record.sha256 == expected_digest,
                f"{record.product_type} digest does not match the stub payload",
            )
            _check(
                record.byte_count == len(payloads[record.product_type]),
                f"{record.product_type} byte count does not match the stub payload",
            )
        checks += 1

        run_record = {
            "task": TASK,
            "status": "OK",
            "mode": "self-test",
            "synthetic": True,
            "tic_id": tic_id,
            "sector": sector,
            "byte_cap_bytes": max_bytes,
            "allowlist": list(ALLOWED_HOSTS),
            "query": query_record,
            "products": product_entries,
            "manifest": _display(manifest),
            "manifest_urls": manifest_urls,
        }
        _write_json(work / "run.json", run_record)
        checks += 1

    _expect_error(
        "no matching product",
        NoMatchingProductError,
        lambda: resolve_products(rows, tic_id=987654321, sector=sector),
    )
    checks += 1

    bad_scheme = dict(rows[0])
    bad_scheme["accessURL"] = (
        "http://mast.stsci.edu/product.fits"  # sec-ok: http-url - deliberate non-https value tests rejection
    )
    _expect_error(
        "non-https URL",
        InsecureUrlError,
        lambda: resolve_products(
            [bad_scheme],
            tic_id=tic_id,
            sector=sector,
            requested_types=("lightcurve",),
        ),
    )
    checks += 1

    bad_host = dict(rows[0])
    bad_host["accessURL"] = "https://example.org/product.fits"
    _expect_error(
        "host outside the allowlist",
        HostNotAllowedError,
        lambda: resolve_products(
            [bad_host],
            tic_id=tic_id,
            sector=sector,
            requested_types=("lightcurve",),
        ),
    )
    checks += 1

    _expect_error(
        "declared byte cap",
        ByteCapExceededError,
        lambda: resolve_products(rows, tic_id=tic_id, sector=sector, max_bytes=10),
    )
    checks += 1

    duplicate = dict(rows[0])
    duplicate["dataURI"] = (
        "mast:TESS/product/tess2026001000000-s0001-000000123456789-0120-s_extra_lc.fits"
    )
    duplicate["productFilename"] = (
        "tess2026001000000-s0001-000000123456789-0120-s_extra_lc.fits"
    )
    duplicate.pop("accessURL", None)
    _expect_error(
        "duplicate product",
        DuplicateProductError,
        lambda: resolve_products(
            [rows[0], duplicate],
            tic_id=tic_id,
            sector=sector,
            requested_types=("lightcurve",),
        ),
    )
    checks += 1

    with tempfile.TemporaryDirectory(
        prefix="fetch-products-digest-self-test-", dir=REPO_ROOT / ".scratch"
    ) as temporary_directory:
        destination = Path(temporary_directory) / "digest-check.fits"
        _expect_error(
            "digest mismatch",
            DigestMismatchError,
            lambda: download_product(
                products[0],
                destination,
                max_bytes=max_bytes,
                opener=stub.open,
                expected_sha256="0" * 64,
            ),
        )
        _check(not destination.exists(), "a rejected digest left a file behind")
        checks += 1

    print("status=OK")
    print("self_test=OK")
    print(f"products={len(products)}")
    print("manifest_records=2")
    print(f"checks={checks}")
    print("downloads=stub")
    print("network=none")
    return 0


def _run_live(args: argparse.Namespace, *, opener: Any | None = None) -> int:
    if args.tic_id is None or args.sector is None:
        raise ProductError(
            "--tic-id and --sector are required unless --self-test is used"
        )
    if args.tic_id <= 0 or args.sector <= 0:
        raise ProductError("--tic-id and --sector must be positive integers")
    if args.max_bytes <= 0:
        raise ByteCapExceededError("--max-bytes must be a positive integer")

    if args.fixture is not None:
        rows, query_record = _fixture_rows(args.fixture)
    else:
        rows, query_record = query_mast_products(
            args.tic_id, args.sector, opener=opener
        )

    requested = tuple(args.product_type or PRODUCT_TYPES)
    products = resolve_products(
        rows,
        tic_id=args.tic_id,
        sector=args.sector,
        requested_types=requested,
        max_bytes=args.max_bytes,
    )

    print("status=OK")
    print(f"tic_id={args.tic_id}")
    print(f"sector={args.sector}")
    print(f"byte_cap_bytes={args.max_bytes}")
    for product in products:
        print(
            f"product={product.product_type} cadence={product.cadence} "
            f"declared_bytes={product.declared_bytes} url={product.file_url}"
        )

    product_records: list[dict[str, Any]] = []
    manifest_count = None
    for product in products:
        record = _product_record(product)
        if args.download:
            cadence = require_cadence(product)
            destination = (
                args.downloads / str(args.tic_id) / str(args.sector) / product.filename
            )
            result = download_product(
                product,
                destination,
                max_bytes=args.max_bytes,
                opener=opener,
            )
            provenance = record_from_file(
                destination,
                target=product.target_name,
                tic_id=args.tic_id,
                sector=args.sector,
                cadence=cadence,
                product_type=product.product_type,
                file_url=product.file_url,
                aperture=args.aperture,
                quality_mask=args.quality_mask,
                code_version=DEFAULT_CODE_VERSION,
                downloaded_utc=result.downloaded_utc,
            )
            manifest_count = append_record(
                args.manifest, provenance, file_path=destination
            )
            record.update(
                {
                    "http_status": result.http_status,
                    "requested_utc": result.requested_utc,
                    "downloaded_utc": result.downloaded_utc,
                    "byte_count": result.byte_count,
                    "sha256": result.sha256,
                    "local_path": _display(destination),
                }
            )
            print(
                f"downloaded={product.product_type} bytes={result.byte_count} "
                f"sha256={result.sha256} path={_display(destination)}"
            )
        product_records.append(record)

    if args.download:
        verify_manifest(args.manifest)
        print(f"manifest={_display(args.manifest)}")
        print(f"manifest_records={manifest_count}")

    run_record = {
        "task": TASK,
        "status": "OK",
        "mode": "fixture" if args.fixture is not None else "live",
        "synthetic": bool(query_record.get("fixture_synthetic", False)),
        "tic_id": args.tic_id,
        "sector": args.sector,
        "byte_cap_bytes": args.max_bytes,
        "allowlist": list(ALLOWED_HOSTS),
        "query": query_record,
        "products": product_records,
        "manifest": _display(args.manifest) if args.download else None,
    }
    _write_json(args.out, run_record)
    print(f"run_record={_display(args.out)}")
    return 0


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Resolve one TESS light-curve product and one target-pixel product "
            "for a TIC identifier and sector. The default is a read-only MAST "
            "metadata query; --download fetches the resolved files and appends "
            "provenance records. The tool never chooses a target."
        )
    )
    parser.add_argument(
        "--tic-id",
        type=int,
        default=None,
        help="TIC identifier (required unless --self-test is used; no default)",
    )
    parser.add_argument(
        "--sector",
        type=int,
        default=None,
        help="TESS sector number (required unless --self-test is used; no default)",
    )
    parser.add_argument(
        "--product-type",
        action="append",
        choices=PRODUCT_TYPES,
        default=None,
        help=(
            "product type to resolve; repeat for both "
            f"(default: {' and '.join(PRODUCT_TYPES)})"
        ),
    )
    parser.add_argument(
        "--download",
        action="store_true",
        help="download the resolved products and append provenance records",
    )
    parser.add_argument(
        "--max-bytes",
        type=int,
        default=DEFAULT_MAX_BYTES,
        help=f"per-file byte cap before writing (default: {DEFAULT_MAX_BYTES})",
    )
    parser.add_argument(
        "--aperture",
        default="not applied",
        help="aperture label recorded in provenance (default: not applied)",
    )
    parser.add_argument(
        "--quality-mask",
        default="not applied",
        help="quality-mask label recorded in provenance (default: not applied)",
    )
    parser.add_argument(
        "--manifest",
        type=Path,
        default=DEFAULT_MANIFEST,
        help=f"provenance JSONL manifest (default: {_display(DEFAULT_MANIFEST)})",
    )
    parser.add_argument(
        "--downloads",
        type=Path,
        default=DEFAULT_DOWNLOADS,
        help=f"download root (default: {_display(DEFAULT_DOWNLOADS)})",
    )
    parser.add_argument(
        "--out",
        type=Path,
        default=DEFAULT_OUT,
        help=f"JSON run record (default: {_display(DEFAULT_OUT)})",
    )
    parser.add_argument(
        "--fixture",
        type=Path,
        default=None,
        help="read product rows from a committed MAST response fixture instead of querying MAST",
    )
    parser.add_argument(
        "--self-test",
        action="store_true",
        help="run the offline self-test with a committed fixture and a stub opener",
    )
    return parser


def main(argv: list[str] | None = None, *, opener: Any | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    try:
        if args.self_test:
            if args.tic_id is not None or args.sector is not None:
                parser.error("--self-test cannot be combined with --tic-id or --sector")
            return _run_self_test()
        return _run_live(args, opener=opener)
    except (ProductError, ProvenanceError, OSError) as exc:
        print("status=FAILED", flush=True)
        print(f"error={type(exc).__name__}: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())