#!/usr/bin/env python3
"""SPTI-shaped JSON/JSONL client for the GDR/XDVD TCP host bridge.

This is not a kernel driver and not an iSCSI target.  It is a user-mode adapter
that accepts request objects shaped like Windows SCSI_PASS_THROUGH/_DIRECT test
records and translates them to the stable GDR JSONL TCP bridge protocol.

Input examples:
  {"tag":"tur", "Cdb":[0,0,0,0,0,0], "DataIn":"none"}
  {"tag":"inq", "CdbHex":"12 00 00 00 24 00", "DataIn":"in", "DataTransferLength":36}
  {"tag":"read10", "CdbHex":"28 00 00 00 00 20 00 00 01 00", "DataIn":"in"}
  {"tag":"force_unlocked", "op":"reset_auth", "stage":"unlocked"}

Output is JSONL with SPTI-style fields plus the full underlying host response in
"host_response" for debugging.
"""
from __future__ import annotations

import argparse, base64, json, socket, sys
from pathlib import Path
from typing import Any

from gdr_host_tcp_client import read_ready_file
from gdr_host_jsonl_bridge import load_requests, parse_hex_bytes


def hx(data: bytes) -> str:
    return ' '.join(f'{b:02X}' for b in data)


def _bytes_from_any(value: Any) -> bytes:
    if value is None:
        return b''
    if isinstance(value, (bytes, bytearray)):
        return bytes(value)
    if isinstance(value, list):
        return bytes((int(x, 0) if isinstance(x, str) else int(x)) & 0xFF for x in value)
    if isinstance(value, dict):
        if 'hex' in value:
            return parse_hex_bytes(value.get('hex'))
        if 'base64' in value:
            return base64.b64decode(str(value.get('base64') or ''))
    return parse_hex_bytes(value)


def cdb_from_spti(req: dict[str, Any]) -> bytes:
    for key in ('CdbHex','CDBHex','cdb_hex','cdb'):
        if req.get(key) is not None:
            return _bytes_from_any(req.get(key))
    for key in ('Cdb','CDB'):
        if req.get(key) is not None:
            return _bytes_from_any(req.get(key))
    return b''


def data_out_from_spti(req: dict[str, Any]) -> bytes:
    for key in ('DataOutHex','data_out_hex','DataOut','data_out'):
        if req.get(key) is not None:
            return _bytes_from_any(req.get(key))
    for key in ('DataOutBase64','data_out_base64'):
        if req.get(key) is not None:
            return base64.b64decode(str(req.get(key) or ''))
    return b''


def to_host_request(req: dict[str, Any], index: int) -> dict[str, Any]:
    out: dict[str, Any] = {'tag': req.get('tag') or req.get('Tag') or f'spti_{index:03d}'}
    # Bridge control ops are allowed so the same SPTI-shaped harness can test
    # tray/media/auth transitions around real CDBs.
    if req.get('op'):
        for k in ('op','stage','xdvd_json','xdvd','iso','iso_path','log','log_path','force','reset_auth','unlocked'):
            if k in req:
                out[k] = req[k]
        return out
    cdb = cdb_from_spti(req)
    if not cdb:
        raise ValueError(f'empty CDB in request index {index}')
    out['cdb_hex'] = hx(cdb)
    dout = data_out_from_spti(req)
    if dout:
        out['data_out_hex'] = hx(dout)
    return out


def json_line(obj: dict[str, Any]) -> bytes:
    return (json.dumps(obj, separators=(',', ':')) + '\n').encode('utf-8')


def send_host_requests(host: str, port: int, host_reqs: list[dict[str, Any]], timeout: float) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with socket.create_connection((host, port), timeout=timeout) as sock:
        sock.settimeout(timeout)
        for req in host_reqs:
            sock.sendall(json_line(req))
        sock.shutdown(socket.SHUT_WR)
        buf = b''
        while True:
            chunk = sock.recv(65536)
            if not chunk:
                break
            buf += chunk
    for line in buf.decode('utf-8', errors='replace').splitlines():
        if line.strip():
            rows.append(json.loads(line))
    if len(rows) != len(host_reqs):
        raise SystemExit(f'SPTI client response count mismatch: got={len(rows)} expected={len(host_reqs)}')
    return rows


def spti_status_value(row: dict[str, Any]) -> int:
    return 0x00 if row.get('scsi_status') == 'GOOD' else 0x02


def _int_field(req: dict[str, Any], names: tuple[str, ...], default: int | None = None) -> int | None:
    for name in names:
        if req.get(name) is not None:
            try:
                return int(req.get(name), 0) if isinstance(req.get(name), str) else int(req.get(name))
            except Exception:
                return default
    return default


def _data_transfer_limit(req: dict[str, Any]) -> int | None:
    # SCSI_PASS_THROUGH/_DIRECT uses DataTransferLength as the host buffer size.
    # Some loggers name the same concept DataBufferLength/TransferLength, so accept
    # those spellings too.  None means no extra host-side clamp beyond the CDB's
    # own allocation/transfer length that the parser already applied.
    return _int_field(req, ('DataTransferLength','data_transfer_length','DataBufferLength','data_buffer_length','TransferLength','transfer_length'), None)


def _sense_transfer_limit(req: dict[str, Any]) -> int | None:
    return _int_field(req, ('SenseInfoLength','sense_info_length','SenseBufferLength','sense_buffer_length'), None)


def convert_response(req: dict[str, Any], host_req: dict[str, Any], row: dict[str, Any], index: int, response_dir: Path | None) -> dict[str, Any]:
    raw_data = parse_hex_bytes(row.get('data_hex','')) if row.get('data_hex') else b''
    raw_sense = parse_hex_bytes(row.get('sense_data_hex','')) if row.get('sense_data_hex') else b''
    data_limit = _data_transfer_limit(req)
    sense_limit = _sense_transfer_limit(req)
    data = raw_data[:max(0, data_limit)] if data_limit is not None else raw_data
    sense = raw_sense[:max(0, sense_limit)] if sense_limit is not None else raw_sense
    data_file = ''
    sense_file = ''
    if response_dir:
        response_dir.mkdir(parents=True, exist_ok=True)
        safe = str(row.get('tag') or f'spti_{index:03d}').replace('/','_').replace('\\','_')
        if data:
            data_file = f'{index:03d}_{safe}_data.bin'
            (response_dir / data_file).write_bytes(data)
        if sense:
            sense_file = f'{index:03d}_{safe}_sense.bin'
            (response_dir / sense_file).write_bytes(sense)
    requested = data_limit if data_limit is not None else len(raw_data)
    sense_requested = sense_limit if sense_limit is not None else len(raw_sense)
    scsi_status = spti_status_value(row)
    check_condition = (row.get('scsi_status') == 'CHECK CONDITION')
    # Windows-style surface for SCSI_PASS_THROUGH/SCSI_PASS_THROUGH_DIRECT
    # harnesses.  This bridge models a successful transport/DeviceIoControl call
    # whose target status may still be CHECK CONDITION.  SRB_STATUS_SUCCESS is
    # 0x01; for check condition with autosense we use 0x84
    # (SRB_STATUS_ERROR | SRB_STATUS_AUTOSENSE_VALID).
    srb_status = 0x84 if check_condition else 0x01
    srb_status_name = 'SRB_STATUS_ERROR_AUTOSENSE_VALID' if check_condition else 'SRB_STATUS_SUCCESS'
    device_io_control_ok = True
    win32_error = 0
    bytes_returned = len(data) + len(sense)
    return {
        'index': index,
        'tag': row.get('tag') or req.get('tag') or f'spti_{index:03d}',
        'ok': bool(row.get('ok')),
        'DeviceIoControlOk': device_io_control_ok,
        'Win32Error': win32_error,
        'BytesReturned': bytes_returned,
        'ScsiStatus': scsi_status,
        'ScsiStatusName': row.get('scsi_status',''),
        'SrbStatus': srb_status,
        'SrbStatusName': srb_status_name,
        'HostAdapterStatus': 0,
        'CdbLength': len(parse_hex_bytes(host_req.get('cdb_hex',''))) if host_req.get('cdb_hex') else 0,
        'TimeOutValue': _int_field(req, ('TimeOutValue','timeout_value','TimeoutValue'), 0) or 0,
        'DataIn': req.get('DataIn', req.get('data_in', '')),
        'SenseInfoLengthRequested': sense_requested,
        'SenseInfoLength': len(sense),
        'SenseInfoLengthRaw': len(raw_sense),
        'SenseInfoResidual': max(0, int(sense_requested or 0) - len(sense)),
        'SenseInfoHex': hx(sense),
        'SenseInfoBase64': base64.b64encode(sense).decode('ascii') if sense else '',
        'SenseInfoHexTruncated': len(raw_sense) > len(sense),
        'SenseInfoFile': sense_file,
        'DataTransferLengthRequested': requested,
        'DataTransferLength': len(data),
        'DataTransferLengthRaw': len(raw_data),
        'DataTransferLengthResidual': max(0, int(requested or 0) - len(data)),
        'DataInHex': hx(data[:256]),
        'DataInBase64': base64.b64encode(data).decode('ascii') if data else '',
        'DataInHexTruncated': len(data) > 256,
        'DataInFile': data_file,
        'CdbHex': host_req.get('cdb_hex',''),
        'Direction': row.get('direction','none'),
        'media_present': row.get('media_present'),
        'tray_state': row.get('tray_state'),
        'prevent_removal': row.get('prevent_removal'),
        'auth_stage_after': row.get('auth_stage_after'),
        'catalog_match': row.get('catalog_match'),
        'host_response': row,
    }


def write_verify(path: Path, rows: list[dict[str, Any]], *, title: str = 'GDR SPTI-shaped passthrough verification') -> None:
    ok=sum(1 for r in rows if r.get('ok'))
    cc=sum(1 for r in rows if r.get('ScsiStatusName') == 'CHECK CONDITION')
    payload=sum(1 for r in rows if int(r.get('DataTransferLength') or 0) > 0)
    catalog=sum(1 for r in rows if r.get('catalog_match'))
    lines=[f'# {title}', '', f'requests={len(rows)}', f'ok={ok}', f'check_condition={cc}', f'payload_responses={payload}', f'catalog_matches={catalog}', '']
    lines.append('| index | tag | status | ok | len | auth | media | tray | catalog |')
    lines.append('|---:|---|---|---|---:|---|---|---|---|')
    for r in rows:
        lines.append(f"| {r.get('index')} | `{r.get('tag','')}` | {r.get('ScsiStatusName','')} | {str(bool(r.get('ok'))).lower()} | {r.get('DataTransferLength',0)} | {r.get('auth_stage_after','')} | {r.get('media_present','')} | {r.get('tray_state','')} | {str(bool(r.get('catalog_match'))).lower()} |")
    path.write_text('\n'.join(lines)+'\n', encoding='utf-8')


def run_client(args: argparse.Namespace) -> int:
    ready_host, ready_port = read_ready_file(args.ready_file)
    host = ready_host or str(args.host)
    port = ready_port if ready_port is not None else int(args.port)
    reqs = load_requests(args.script)
    host_reqs = [to_host_request(req, i) for i, req in enumerate(reqs)]
    host_rows = send_host_requests(host, port, host_reqs, float(args.timeout))
    rd = Path(args.write_response_dir) if args.write_response_dir else None
    rows = [convert_response(reqs[i], host_reqs[i], host_rows[i], i, rd) for i in range(len(reqs))]
    lines='\n'.join(json.dumps(r, separators=(',',':')) for r in rows) + ('\n' if rows else '')
    if args.out and args.out != '-':
        Path(args.out).write_text(lines, encoding='utf-8')
    else:
        sys.stdout.write(lines)
    if args.verify_md:
        write_verify(Path(args.verify_md), rows)
    print(f'[SPTI-CLIENT] requests={len(rows)} responses={len(rows)} ok={sum(1 for r in rows if r.get("ok"))}', flush=True)
    return 0 if all(r.get('ok') for r in rows) else (0 if args.allow_check_condition else 2)


def main() -> int:
    ap=argparse.ArgumentParser(description='SPTI-shaped client for GDR/XDVD TCP host bridge')
    ap.add_argument('--host', default='127.0.0.1')
    ap.add_argument('--port', type=int, default=32610)
    ap.add_argument('--ready-file')
    ap.add_argument('--timeout', default='10')
    ap.add_argument('--script', '-i', help='SPTI-shaped JSON/JSONL requests; stdin when omitted or -')
    ap.add_argument('--out', '-o', help='write SPTI-shaped JSONL responses; stdout when omitted or -')
    ap.add_argument('--write-response-dir')
    ap.add_argument('--verify-md')
    ap.add_argument('--allow-check-condition', action='store_true')
    args=ap.parse_args()
    return run_client(args)

if __name__ == '__main__':
    raise SystemExit(main())
