#!/usr/bin/env python3
"""
DICOM Sender (C-STORE client)
==============================
Sends DICOM files to a remote DICOM router/PACS over the network using the
C-STORE service, selecting which files to send by matching Accession Number
against an index CSV (the kind produced by dicom_indexer.py's 'scan' command,
which has FilePath and AccessionNumber columns).

This tool always works from TWO separate source files:
  --source-index       the file/metadata index (built earlier by dicom_indexer.py)
  --source-accessions  your list of accession numbers to send (a different file)

Dependencies:
    pip install pydicom pynetdicom

Usage:
    # Test connectivity to the destination router first (recommended)
    python dicom_sender.py echo --host 192.168.1.50 --port 104 --ae-title REMOTE_PACS

    # Send files whose AccessionNumber matches a list in a separate CSV
    python dicom_sender.py send \\
        --source-index dicom_index.csv \\
        --source-accessions accessions_to_send.csv \\
        --host 192.168.1.50 --port 104 --ae-title REMOTE_PACS

    # Send files for a single accession number directly (no accession-list file needed)
    python dicom_sender.py send \\
        --source-index dicom_index.csv \\
        --accession ACC12345 \\
        --host 192.168.1.50 --port 104 --ae-title REMOTE_PACS

    # Use a custom calling (your own) AE Title instead of the default
    python dicom_sender.py send ... --calling-ae-title MY_AE

    # Dry run: show what WOULD be sent without opening a network connection
    python dicom_sender.py send ... --dry-run
"""

import os
import sys
import csv
import argparse
import socket
import time

DEFAULT_CALLING_AE_TITLE = "DCM-RSABHK"

# ---------------------------------------------------------------------------
# Dependency check with a friendly message
# ---------------------------------------------------------------------------
MISSING = []
try:
    import pydicom
except ImportError:
    MISSING.append("pydicom")

try:
    from pynetdicom import AE, evt
    from pynetdicom.sop_class import Verification
except ImportError:
    MISSING.append("pynetdicom")

if MISSING:
    print("Missing required packages: " + ", ".join(MISSING))
    print("Install them with:")
    print(f"    pip install {' '.join(MISSING)}")
    sys.exit(1)


# ---------------------------------------------------------------------------
# DICOM C-STORE status code meanings (the common ones; peers may return
# vendor-specific codes outside this list, which we just show as hex).
# ---------------------------------------------------------------------------
STATUS_MEANINGS = {
    0x0000: "Success",
    0x0105: "No such attribute",
    0x0106: "Invalid attribute value",
    0x0110: "Processing failure",
    0x0117: "Invalid SOP instance",
    0x0122: "SOP class not supported",
    0x0124: "Not authorised",
    0x0210: "Duplicate invocation",
    0x0211: "Unrecognised operation",
    0x0212: "Mistyped argument",
    0xA700: "Out of resources",
    0xA900: "Data set does not match SOP class",
    0xB000: "Coercion of data elements (warning)",
    0xB006: "Element discarded (warning)",
    0xB007: "Data set does not match SOP class (warning)",
}


def describe_status(status_dataset):
    """Return a (code, label) tuple describing a C-STORE/C-ECHO response status."""
    if status_dataset is None or "Status" not in status_dataset:
        return None, "No response (association/timeout/abort failure)"
    code = status_dataset.Status
    label = STATUS_MEANINGS.get(code)
    if label is None:
        # Out-of-resources and "cannot understand" are ranges, not single codes.
        if 0xA700 <= code <= 0xA7FF:
            label = "Out of resources"
        elif 0xC000 <= code <= 0xCFFF:
            label = "Cannot understand"
        else:
            label = "Unknown/vendor-specific status"
    return code, label


def is_success(code):
    return code == 0x0000


# ---------------------------------------------------------------------------
# Index / accession-list loading (streamed — see dicom_indexer.py's memory
# notes; we apply the same approach here so large indexes don't balloon RAM)
# ---------------------------------------------------------------------------
def load_accession_list(path):
    """
    Read a CSV of accession numbers to send. Accepts either:
      - a column literally named 'AccessionNumber' (case-insensitive), or
      - a single-column CSV with no recognizable header (first column used).
    Returns a set of accession number strings (stripped, deduplicated).
    """
    accessions = set()
    with open(path, "r", newline="", encoding="utf-8") as f:
        reader = csv.reader(f)
        rows = list(reader)

    if not rows:
        return accessions

    header = [h.strip().lower() for h in rows[0]]
    col_idx = None
    for i, h in enumerate(header):
        if h in ("accessionnumber", "accession_number", "accession"):
            col_idx = i
            break

    if col_idx is not None:
        data_rows = rows[1:]
    else:
        # No recognizable header — treat every row's first column as data,
        # including row 0 (it wasn't actually a header).
        col_idx = 0
        data_rows = rows

    for row in data_rows:
        if len(row) > col_idx:
            value = row[col_idx].strip()
            if value:
                accessions.add(value)

    return accessions


def find_matching_files(index_path, accession_numbers):
    """
    Stream the index CSV and yield (filepath, accession_number, sop_class_uid,
    patient_name, series_description) for every row whose AccessionNumber is
    in accession_numbers.

    accession_numbers: a set of accession number strings to match against.
    """
    with open(index_path, "r", newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        if "FilePath" not in reader.fieldnames or "AccessionNumber" not in reader.fieldnames:
            print(
                "Error: index file does not have the expected 'FilePath' and "
                "'AccessionNumber' columns. Make sure this is an index produced "
                "by dicom_indexer.py's 'scan' command.",
                file=sys.stderr,
            )
            sys.exit(1)

        for row in reader:
            acc = (row.get("AccessionNumber") or "").strip()
            if acc in accession_numbers:
                yield row


# ---------------------------------------------------------------------------
# Subcommand: echo (connectivity test)
# ---------------------------------------------------------------------------
def cmd_echo(args):
    ae = AE(ae_title=args.calling_ae_title)
    ae.add_requested_context(Verification)
    ae.acse_timeout = args.timeout
    ae.dimse_timeout = args.timeout
    ae.network_timeout = args.timeout

    print(
        f"Sending C-ECHO to {args.ae_title}@{args.host}:{args.port} "
        f"(calling AE: {args.calling_ae_title})..."
    )

    try:
        assoc = ae.associate(args.host, args.port, ae_title=args.ae_title)
    except Exception as exc:
        print(f"Association error: {exc}", file=sys.stderr)
        sys.exit(1)

    if not assoc.is_established:
        print(
            "FAILED: could not establish association. Check host/port/AE title "
            "and that the destination router is reachable and accepting "
            "associations from this AE.",
            file=sys.stderr,
        )
        sys.exit(1)

    status = assoc.send_c_echo()
    code, label = describe_status(status)
    assoc.release()

    if code is not None and is_success(code):
        print(f"SUCCESS: C-ECHO accepted (status 0x{code:04X} - {label}).")
    else:
        code_str = f"0x{code:04X}" if code is not None else "N/A"
        print(f"FAILED: C-ECHO returned status {code_str} - {label}.", file=sys.stderr)
        sys.exit(1)


# ---------------------------------------------------------------------------
# Subcommand: send
# ---------------------------------------------------------------------------
def cmd_send(args):
    # ---- Step 1: resolve the set of accession numbers to send ----
    # Two independent source files are used here:
    #   args.source_index       -> which files exist + their metadata
    #   args.source_accessions  -> which accession numbers to actually send
    accession_numbers = set()
    if args.source_accessions:
        if not os.path.isfile(args.source_accessions):
            print(f"Error: accession source file not found: {args.source_accessions}", file=sys.stderr)
            sys.exit(1)
        accession_numbers |= load_accession_list(args.source_accessions)
    if args.accession:
        accession_numbers |= {a.strip() for a in args.accession if a.strip()}

    if not accession_numbers:
        print(
            "Error: no accession numbers given. Use --source-accessions <csv> and/or "
            "one or more --accession <value>.",
            file=sys.stderr,
        )
        sys.exit(1)

    if not os.path.isfile(args.source_index):
        print(f"Error: index source file not found: {args.source_index}", file=sys.stderr)
        sys.exit(1)

    print(f"Looking up {len(accession_numbers)} accession number(s) in index '{args.source_index}'...")

    # ---- Step 2: find matching files in the index ----
    matched_rows = list(find_matching_files(args.source_index, accession_numbers))

    if not matched_rows:
        print("No matching files found in the index for the given accession number(s).")
        sys.exit(0)

    found_accessions = {r["AccessionNumber"] for r in matched_rows}
    missing_accessions = accession_numbers - found_accessions
    print(f"Found {len(matched_rows)} file(s) across {len(found_accessions)} accession number(s).")
    if missing_accessions:
        print(
            f"Warning: {len(missing_accessions)} accession number(s) had no matching "
            f"file in the index: {', '.join(sorted(missing_accessions)[:10])}"
            + (" ..." if len(missing_accessions) > 10 else ""),
            file=sys.stderr,
        )

    # ---- Step 3: verify files exist on disk, and (for real sends) preload
    #      each file's SOP Class UID + Transfer Syntax UID so we can build
    #      the minimal set of presentation contexts to request ----
    sendable = []
    missing_files = []
    for row in matched_rows:
        filepath = row["FilePath"]
        if os.path.isfile(filepath):
            sendable.append(row)
        else:
            missing_files.append(filepath)

    if missing_files:
        print(
            f"Warning: {len(missing_files)} file(s) listed in the index no longer "
            f"exist on disk and will be skipped:",
            file=sys.stderr,
        )
        for fp in missing_files[:10]:
            print(f"  - {fp}", file=sys.stderr)
        if len(missing_files) > 10:
            print(f"  ... and {len(missing_files) - 10} more", file=sys.stderr)

    if not sendable:
        print("No sendable files remain after checking disk. Aborting.")
        sys.exit(1)

    if args.dry_run:
        print("\n--- DRY RUN: the following files would be sent ---")
        for row in sendable:
            print(f"  [{row.get('AccessionNumber','')}] {row['FilePath']}")
        print(f"\nTotal: {len(sendable)} file(s). No network connection was made.")
        return

    # ---- Step 4: read each file's SOP Class UID / Transfer Syntax UID to
    #      build presentation contexts (max 128 per association) ----
    print("Reading SOP Class / Transfer Syntax info from files to build presentation contexts...")
    file_info = []  # (filepath, row, sop_class_uid, transfer_syntax_uid) or None on read failure
    unreadable = []
    for row in sendable:
        filepath = row["FilePath"]
        try:
            ds_meta = pydicom.dcmread(filepath, stop_before_pixels=True, force=True)
            sop_class = getattr(ds_meta, "SOPClassUID", None)
            ts = ds_meta.file_meta.TransferSyntaxUID if hasattr(ds_meta, "file_meta") else None
            if sop_class is None:
                unreadable.append(filepath)
                continue
            file_info.append((filepath, row, str(sop_class), str(ts) if ts else None))
        except Exception as exc:
            unreadable.append(filepath)

    if unreadable:
        print(
            f"Warning: {len(unreadable)} file(s) could not be read as DICOM and "
            f"will be skipped:",
            file=sys.stderr,
        )
        for fp in unreadable[:10]:
            print(f"  - {fp}", file=sys.stderr)

    if not file_info:
        print("No readable DICOM files remain. Aborting.")
        sys.exit(1)

    # Build the minimal set of presentation contexts needed: one context per
    # DISTINCT (SOP Class UID, Transfer Syntax UID) pair actually used by the
    # files being sent. We deliberately do NOT bundle multiple transfer
    # syntaxes into a single context, because the peer only accepts ONE
    # transfer syntax per context — if a SOP Class has some files in (say)
    # uncompressed Explicit VR LE and others in JPEG Lossless, bundling them
    # into one context risks the peer accepting only one of the two, and
    # files using the other transfer syntax would still fail to send. A
    # separate context per pair guarantees every transfer syntax actually
    # present gets its own negotiation slot.
    #
    # This also fixes the root cause of "No presentation context ... has
    # been accepted by the peer with '<transfer syntax>'" errors: previously
    # contexts were requested using pynetdicom's DEFAULT transfer syntax
    # list (uncompressed only), so compressed files (JPEG Lossless, JPEG
    # 2000, RLE, etc.) never had a matching context at all.
    seen_pairs = set()
    context_pairs = []  # list of (sop_class, transfer_syntax_or_None)
    for _, _, sop_class, ts in file_info:
        pair = (sop_class, ts)
        if pair not in seen_pairs:
            seen_pairs.add(pair)
            context_pairs.append(pair)

    if len(context_pairs) > 128:
        print(
            f"Error: files span {len(context_pairs)} distinct SOP "
            f"Class/Transfer Syntax combinations, which exceeds the 128 "
            f"presentation contexts allowed per association. Split the "
            f"send into smaller batches.",
            file=sys.stderr,
        )
        sys.exit(1)

    # ---- Step 5: build AE, request contexts, associate ----
    ae = AE(ae_title=args.calling_ae_title)
    for sop_class, ts in context_pairs:
        if ts:
            ae.add_requested_context(sop_class, ts)
        else:
            # Transfer syntax couldn't be determined for this file
            # (unexpected/corrupt file_meta) — fall back to pynetdicom's
            # default uncompressed transfer syntax list rather than failing.
            ae.add_requested_context(sop_class)

    ae.acse_timeout = args.timeout
    ae.dimse_timeout = args.timeout
    ae.network_timeout = args.timeout

    distinct_sop_classes = {p[0] for p in context_pairs}
    print(
        f"\nConnecting to {args.ae_title}@{args.host}:{args.port} "
        f"(calling AE: {args.calling_ae_title}) requesting "
        f"{len(context_pairs)} presentation context(s) covering "
        f"{len(distinct_sop_classes)} SOP Class(es)..."
    )

    try:
        assoc = ae.associate(args.host, args.port, ae_title=args.ae_title)
    except Exception as exc:
        print(f"Association error: {exc}", file=sys.stderr)
        sys.exit(1)

    if not assoc.is_established:
        print(
            "FAILED: could not establish association. Check host/port/AE title, "
            "that the destination accepts the SOP classes being sent, and that "
            "it allows associations from this calling AE title.",
            file=sys.stderr,
        )
        sys.exit(1)

    # ---- Step 6: send each file via C-STORE ----
    results = []  # (filepath, accession, code, label)
    try:
        for i, (filepath, row, sop_class, ts) in enumerate(file_info, start=1):
            accession = row.get("AccessionNumber", "")
            sys.stdout.write(
                f"\rSending {i}/{len(file_info)} (accession {accession})..."
            )
            sys.stdout.flush()

            try:
                status = assoc.send_c_store(filepath)
            except Exception as exc:
                results.append((filepath, accession, None, f"Exception: {exc}"))
                continue

            code, label = describe_status(status)
            results.append((filepath, accession, code, label))

            if args.delay:
                time.sleep(args.delay)
    finally:
        assoc.release()

    print()  # newline after progress line

    # ---- Step 7: report results ----
    succeeded = [r for r in results if r[2] is not None and is_success(r[2])]
    failed = [r for r in results if r not in succeeded]

    print(f"\n--- Send Summary ---")
    print(f"Total attempted: {len(results)}")
    print(f"Succeeded:       {len(succeeded)}")
    print(f"Failed:          {len(failed)}")

    if failed:
        print("\nFailed files:")
        for filepath, accession, code, label in failed:
            code_str = f"0x{code:04X}" if code is not None else "N/A"
            print(f"  [{accession}] {filepath}")
            print(f"      -> status {code_str}: {label}")

    if args.report:
        with open(args.report, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow(["FilePath", "AccessionNumber", "StatusCode", "StatusLabel"])
            for filepath, accession, code, label in results:
                code_str = f"0x{code:04X}" if code is not None else ""
                writer.writerow([filepath, accession, code_str, label])
        print(f"\nDetailed report saved to: {args.report}")

    if failed:
        sys.exit(1)


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------
def build_parser():
    parser = argparse.ArgumentParser(
        prog="dicom_sender.py",
        description=(
            "Send DICOM files to a remote DICOM router/PACS via C-STORE, "
            "selecting files by Accession Number from a dicom_indexer.py index."
        ),
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Shared connection arguments
    def add_connection_args(p):
        p.add_argument("--host", required=True, help="Destination IP address or hostname.")
        p.add_argument("--port", required=True, type=int, help="Destination port number.")
        p.add_argument(
            "--ae-title", required=True, dest="ae_title",
            help="Destination (called) AE Title.",
        )
        p.add_argument(
            "--calling-ae-title", dest="calling_ae_title", default=DEFAULT_CALLING_AE_TITLE,
            help=f"Your own (calling) AE Title (default: {DEFAULT_CALLING_AE_TITLE}).",
        )
        p.add_argument(
            "--timeout", type=int, default=30,
            help="Network/association/DIMSE timeout in seconds (default: 30).",
        )

    # echo
    p_echo = subparsers.add_parser("echo", help="Test connectivity to a destination via C-ECHO.")
    add_connection_args(p_echo)
    p_echo.set_defaults(func=cmd_echo)

    # send
    p_send = subparsers.add_parser("send", help="Send DICOM files matching accession number(s) via C-STORE.")
    add_connection_args(p_send)
    p_send.add_argument(
        "--source-index", "--index", dest="source_index", required=True,
        help=(
            "SOURCE FILE 1: path to the index CSV produced by dicom_indexer.py's "
            "'scan' command (needs FilePath + AccessionNumber columns). This is "
            "the catalog of files that exist and where they are on disk."
        ),
    )
    p_send.add_argument(
        "--source-accessions", "--accessions", dest="source_accessions",
        help=(
            "SOURCE FILE 2: path to a separate CSV file containing the accession "
            "numbers you want to send (column named AccessionNumber, or a plain "
            "single-column list with no header). This is the work order — which "
            "of the files cataloged in --source-index should actually be sent."
        ),
    )
    p_send.add_argument(
        "--accession", action="append",
        help="A single accession number to send. Can be given multiple times.",
    )
    p_send.add_argument(
        "--dry-run", action="store_true",
        help="Show which files would be sent without making any network connection.",
    )
    p_send.add_argument(
        "--delay", type=float, default=0.0,
        help="Delay in seconds between sending each file (default: 0, no delay).",
    )
    p_send.add_argument(
        "-o", "--report",
        help="Save a detailed per-file CSV report (FilePath, AccessionNumber, StatusCode, StatusLabel) to this path.",
    )
    p_send.set_defaults(func=cmd_send)

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()