#!/usr/bin/env python3
"""v131 stable persistent GDR-8050L firmware session baseline.

Boots the firmware once and feeds a sequential command queue without resetting
between commands.  By default the queue is parsed from the captured GDR unlock
log, but --command-json may be used for a generic command queue.

The D108 hardware data-phase completion model is enabled by default.  v130 A/B
runs showed that this model is required for a complete persistent unlock-log
session, so v131 treats it as baseline behavior while retaining
--no-d108-hw-data-phase as a regression/debug flag.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Any

WORKSPACE = os.environ.get('XDVD_WORKSPACE', r'D:\vscode-workspace\xbox-dvd-tools-v2')
FWBIN = os.environ.get('GDR8050L_FWBIN', r'D:\vscode-workspace\xbox-dvd-tools-v2\firmware\GDR-8050L\8050L_fw_plaintext_256k.bin')
SECTION_RE = re.compile(r"^--- (?P<kind>Sending CDB|Received Data|Sending Data) \[(?P<label>[^\]]+)\] \((?P<n>\d+) bytes\) ---")
HEX_RE = re.compile(r"\b[0-9A-Fa-f]{2}\b")


def parse_hex_bytes(text: Any) -> bytes:
    if text is None:
        return b''
    if isinstance(text, (bytes, bytearray)):
        return bytes(text)
    if isinstance(text, list):
        return bytes((int(x, 0) if isinstance(x, str) else int(x)) & 0xFF for x in text)
    toks = str(text).replace(',', ' ').split()
    out = []
    for tok in toks:
        if tok == '...':
            break
        out.append(int(tok, 16) & 0xFF)
    return bytes(out)


def parse_log(path: Path) -> dict[tuple[str, str], bytes]:
    out: dict[tuple[str, str], bytes] = {}
    lines = path.read_text(encoding='utf-8', errors='replace').splitlines()
    i = 0
    while i < len(lines):
        m = SECTION_RE.match(lines[i].strip())
        if not m:
            i += 1
            continue
        kind, label, n = m.group('kind'), m.group('label'), int(m.group('n'))
        i += 1
        toks: list[str] = []
        while i < len(lines) and not lines[i].startswith('--- '):
            if lines[i].startswith('----------------'):
                i += 1
                break
            toks.extend(HEX_RE.findall(lines[i]))
            i += 1
        data = bytes(int(x, 16) for x in toks[:n])
        if len(data) == n:
            out[(kind, label)] = data
    return out


def require(sections: dict[tuple[str, str], bytes], kind: str, label: str) -> bytes:
    key = (kind, label)
    if key not in sections:
        raise SystemExit(f'[-] missing {kind} [{label}] in unlock log')
    return sections[key]


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


def commands_from_unlock_log(log_path: Path) -> list[dict[str, str]]:
    sec = parse_log(log_path)
    steps = [
        ('step01_locked_read_capacity', 'in',  'Step 1: Init Read Cap',            'Step 1: Init Read Cap'),
        ('step02_locked_mode_sense_3e', 'in',  'Step 2: Mode Sense',               'Step 2: Mode Sense'),
        ('step03_read_dvd_structure_ad_c0', 'in', 'Step 3: Read DVD Struct 0xAD',  'Step 3: Read DVD Struct 0xAD'),
        ('step04_mode_select_chal1', 'out',    'Step 4: Mode Select Chal 1',       'Step 4: Mode Select Chal 1'),
        ('step05_mode_sense_verify1', 'in',    'Step 5: Mode Sense Verify 1',      'Step 5: Mode Sense Verify 1'),
        ('step06_mode_select_chal2', 'out',    'Step 6: Mode Select Chal 2',       'Step 6: Mode Select Chal 2'),
        ('step07_mode_sense_verify2', 'in',    'Step 7: Mode Sense Verify 2',      'Step 7: Mode Sense Verify 2'),
        ('step08_mode_select_partition1_unlock', 'out', 'Step 8: Mode Select Part1 Unlock', 'Step 8: Mode Select Part1 Unlock'),
        ('step09_sticky_descramble_mode31', 'out', 'Step 9: Sticky Descrambling',  'Step 9: Sticky Descrambling'),
        ('step10_unlocked_read_capacity', 'in', 'Step 10: Final Verify',           'Step 10: Final Verify'),
    ]
    commands: list[dict[str, str]] = []
    for tag, direction, cdb_label, data_label in steps:
        cdb = require(sec, 'Sending CDB', cdb_label)
        data = require(sec, 'Received Data' if direction == 'in' else 'Sending Data', data_label)
        entry: dict[str, str] = {'tag': tag, 'title': tag, 'direction': direction, 'cdb_hex': hx(cdb)}
        if direction == 'in':
            entry['data_in_hex'] = hx(data)
        else:
            entry['data_out_hex'] = hx(data)
        commands.append(entry)
    return commands


def commands_from_json(path: Path) -> list[dict[str, str]]:
    raw = json.loads(path.read_text(encoding='utf-8'))
    items = raw.get('commands', raw) if isinstance(raw, dict) else raw
    if not isinstance(items, list):
        raise SystemExit('[-] --command-json must contain either a list or an object with a commands list')
    commands: list[dict[str, str]] = []
    for i, item in enumerate(items):
        if not isinstance(item, dict):
            raise SystemExit(f'[-] command {i} is not an object')
        direction = str(item.get('direction', '')).lower()
        if direction not in {'in', 'out', 'none'}:
            raise SystemExit(f'[-] command {i} has invalid direction {direction!r}; expected in/out/none')
        cdb = parse_hex_bytes(item.get('cdb_hex', item.get('cdb')))
        if not cdb:
            raise SystemExit(f'[-] command {i} has no CDB')
        tag = str(item.get('tag') or item.get('title') or f'cmd{i+1:02d}_{cdb[0]:02x}')
        entry: dict[str, str] = {
            'tag': tag,
            'title': str(item.get('title') or tag),
            'direction': direction,
            'cdb_hex': hx(cdb),
        }
        data_in = parse_hex_bytes(item.get('data_in_hex', item.get('data_in')))
        data_out = parse_hex_bytes(item.get('data_out_hex', item.get('data_out')))
        if direction == 'in':
            entry['data_in_hex'] = hx(data_in)
        elif direction == 'out':
            entry['data_out_hex'] = hx(data_out)
        commands.append(entry)
    return commands


def run_checked(cmd: list[str], cwd: Path, stdout_file: Path | None = None) -> None:
    if stdout_file:
        with stdout_file.open('w', encoding='utf-8', errors='replace') as fh:
            res = subprocess.run(cmd, cwd=str(cwd), stdout=fh, stderr=subprocess.STDOUT)
    else:
        res = subprocess.run(cmd, cwd=str(cwd))
    if res.returncode != 0:
        raise SystemExit(res.returncode)


def minimal_watchpoints() -> list[dict[str, str]]:
    return [
        {'start': '0x0000DF70', 'size': '0x40', 'kind': 'write', 'name': 'watch_df70'},
        {'start': '0x0000DF90', 'size': '0x08', 'kind': 'read',  'name': 'watch_df90'},
        {'start': '0x0000DF2A', 'size': '0x01', 'kind': 'write', 'name': 'watch_df2a'},
        {'start': '0x0000D108', 'size': '0x02', 'kind': 'access','name': 'watch_D108_data_phase'},
        {'start': '0x000005CD', 'size': '0x02', 'kind': 'access','name': 'watch_05CD_auth_candidate'},
    ]


def debug_watchpoints() -> list[dict[str, str]]:
    return minimal_watchpoints() + [
        {'start': '0x000005B1', 'size': '0x01', 'kind': 'access', 'name': 'watch_05B1'},
        {'start': '0x00000143', 'size': '0x01', 'kind': 'access', 'name': 'watch_0143'},
        {'start': '0x0000014A', 'size': '0x01', 'kind': 'access', 'name': 'watch_014A'},
        {'start': '0x000000DE', 'size': '0x01', 'kind': 'access', 'name': 'watch_00DE_host_gate'},
        {'start': '0x0000016D', 'size': '0x01', 'kind': 'access', 'name': 'watch_016D_host_gate'},
        {'start': '0x000008A8', 'size': '0x08', 'kind': 'access', 'name': 'watch_08A8_service_state'},
        {'start': '0x000007AC', 'size': '0x04', 'kind': 'access', 'name': 'watch_07AC_packet_state'},
        {'start': '0x00000778', 'size': '0x01', 'kind': 'access', 'name': 'watch_0778_dispatch'},
        {'start': '0x000007E8', 'size': '0x08', 'kind': 'access', 'name': 'watch_07E8_cdb_mirror'},
        {'start': '0x0000D950', 'size': '0x08', 'kind': 'access', 'name': 'watch_D950_packet_status'},
        {'start': '0x0000DA18', 'size': '0x01', 'kind': 'access', 'name': 'watch_DA18_packet_branch'},
        {'start': '0x00000278', 'size': '0x01', 'kind': 'access', 'name': 'watch_0278_media_branch'},
        {'start': '0x000005B6', 'size': '0x02', 'kind': 'access', 'name': 'watch_05B6_scheduler'},
        {'start': '0x000006F1', 'size': '0x03', 'kind': 'access', 'name': 'watch_06F1_dispatch'},
        {'start': '0x00000132', 'size': '0x04', 'kind': 'access', 'name': 'watch_0132_host_flags'},
        {'start': '0x0000002C', 'size': '0x02', 'kind': 'read',   'name': 'watch_002C_service_irq'},
        {'start': '0x0000D904', 'size': '0x02', 'kind': 'access', 'name': 'watch_D904_packet_status'},
    ]


def minimal_breakpoints() -> list[dict[str, str]]:
    return [
        {'address': '0x90006608', 'name': 'main_loop_90006608'},
        {'address': '0x90008824', 'name': 'packet_timer_done_90008824'},
        {'address': '0x9001CD4D', 'name': 'data_phase_helper_9001CD4D'},
    ]


def debug_breakpoints() -> list[dict[str, str]]:
    return minimal_breakpoints() + [
        {'address': '0x9001BBAC', 'name': 'host_gate_9001BBAC'},
        {'address': '0x900240EC', 'name': 'packet_service_gate_900240EC'},
        {'address': '0x900210C4', 'name': 'packet_service_900210C4'},
        {'address': '0x90024283', 'name': 'packet_fe_frame_copy_90024283'},
        {'address': '0x90021169', 'name': 'packet_tail_copy_90021169'},
        {'address': '0x900222C9', 'name': 'data_phase_dispatch_900222C9'},
        {'address': '0x900222FF', 'name': 'data_phase_ready_900222FF'},
        {'address': '0x90022301', 'name': 'data_phase_call_9001CD4D_90022301'},
        {'address': '0x900099EA', 'name': 'gdr_data_transfer_helper_900099EA'},
        {'address': '0x900099FF', 'name': 'gdr_data_in_source_read_900099FF'},
        {'address': '0x90009A02', 'name': 'gdr_df70_write_90009A02'},
        {'address': '0x90009A0E', 'name': 'gdr_df2a_submit_90009A0E'},
        {'address': '0x90017775', 'name': 'irq_status_handler_90017775'},
        {'address': '0x9001B285', 'name': 'main_scheduler_9001B285'},
        {'address': '0x9001B71B', 'name': 'packet_scheduler_9001B71B'},
    ]


def build_config(args: argparse.Namespace, commands: list[dict[str, str]], stem: str) -> dict[str, Any]:
    trace_csv = f'{stem}_trace.csv'
    debug = bool(args.debug_trace)
    cfg: dict[str, Any] = {
        'test_name': 'GDR-8050L v131 persistent session baseline',
        'workspace': WORKSPACE,
        'target_profile': 'GDR-8050L',
        'custom_bin': FWBIN,
        'custom_base_addr': '0x90000000',
        'entry_point': '0x90000020',
        'max_cycles': int(str(args.max_cycles), 0),
        'header_skip_bytes': '0x0',
        'config_version': 'v100',
        'requires_emulator_version': 'V7-direct-config-trace-v100-gcc-version-sync',
        'initial_registers': {'SP': '0x003FF000'},
        'memory_aliases': [{'name': 'SRAM alias 0x80000000->0x00000000', 'from': '0x80000000', 'to': '0x00000000', 'size': '0x400000'}],
        'trace': {
            'enabled': True,
            'live_instruction_log': bool(args.debug_instructions),
            'out_csv': trace_csv,
            'watchpoints': debug_watchpoints() if debug else minimal_watchpoints(),
            'breakpoints': debug_breakpoints() if debug else minimal_breakpoints(),
            'gdr_hardware_model': {
                'enabled': True,
                'emit_events': debug,
                'timer': {'enabled': True, 'addr': '0x0000DC54', 'base': '0xF000', 'ticks_per_cycle': '0x2000'},
                'packet_status': {'enabled': True, 'source': 'v131 persistent session shared GDR D904/D905 model'},
            },
            'gdr_packet_session_model': {
                'enabled': True,
                'emit_events': True,
                'commands': commands,
                'packet_flags': '0x20',
                'main_loop_pc': '0x90006608',
                'timer_done_pc': '0x90008824',
                'stop_on_complete': True,
                'host_event_gate': {
                    'enabled': not bool(args.no_host_gate),
                    'force_016d_active': True,
                    'force_016d_value': '0x01',
                    'force_00de_pollable': True,
                    'force_00de_value': '0x01',
                    'force_08a8': str(args.host_gate_08a8),
                    'force_d950_value': '0x10',
                    'post_frame_d950_value': '0x40',
                    'clear_08a8_after_frame': not bool(args.no_clear_08a8_after_frame),
                    'clear_08a8_after_tail': not bool(args.no_clear_08a8_after_tail),
                    'clear_d950_after_tail': not bool(args.no_clear_d950_after_tail),
                    'd950_zero_after_post_frame_reads': str(args.d950_zero_after_post_frame_reads),
                    'irq_after_d108_data_phase': not bool(args.no_irq_after_d108_data_phase),
                    'complete_data_phase_on_d108': not bool(args.no_d108_hw_data_phase),
                    'seed_mem8': {
                        '0x000005B1': '0x20',
                        '0x000000DE': '0x01',
                        '0x0000016D': '0x01',
                        '0x000008A8': str(args.host_gate_08a8),
                        '0x000008A9': '0x00',
                        '0x0000DA18': '0x00',
                    },
                },
            },
        },
        'notes': [
            'v131: stable persistent session baseline. D108 hardware data-phase completion is the default model.',
            'Use --command-json for generic command queues; use --no-d108-hw-data-phase only as a regression/debug comparison.',
        ],
    }
    if args.iso:
        cfg['media'] = {'type': 'iso9660/raw2048', 'iso_path': args.iso, 'sector_size': 2048, 'present': True, 'tray_state': 'closed', 'source': '--iso'}
    return cfg


def main() -> int:
    ap = argparse.ArgumentParser(description='Run v131 stable persistent GDR-8050L firmware session baseline')
    ap.add_argument('xdvd_json', help='XDVD sidecar path; used to locate default .log and future media metadata')
    ap.add_argument('iso', nargs='?', help='ISO/raw2048 path stored into media metadata')
    ap.add_argument('--log', help='Unlock log path; default is .log next to xdvd_json')
    ap.add_argument('--command-json', help='Generic persistent command queue JSON; overrides unlock-log parsing')
    ap.add_argument('--stem', default='gdr8050l_auth_v131_persistent_session_baseline', help='output stem')
    ap.add_argument('--max-cycles', default='220000')
    ap.add_argument('--no-host-gate', action='store_true', help='disable host-command re-entry gate')
    ap.add_argument('--host-gate-08a8', default='0x04', help='forced 0x08A8 state at 0x900240EC when host gate is enabled')
    ap.add_argument('--no-clear-08a8-after-frame', action='store_true', help='disable 0x08A8 clear edge after packet frame acceptance')
    ap.add_argument('--no-clear-08a8-after-tail', action='store_true', help='do not clear 0x08A8 after packet tail acceptance')
    ap.add_argument('--no-clear-d950-after-tail', action='store_true', help='do not clear D950 after packet tail acceptance')
    ap.add_argument('--d950-zero-after-post-frame-reads', default='4', help='post-frame D950 reads that return 0x40 before decay; -1 means never decay')
    ap.add_argument('--no-irq-after-d108-data-phase', action='store_true', help='disable service IRQ edge after D108 data-phase request')
    ap.add_argument('--no-d108-hw-data-phase', action='store_true', help='regression flag: disable D108 hardware data-phase completion model')
    ap.add_argument('--debug-trace', action='store_true', help='add detailed watchpoints/breakpoints and hardware events')
    ap.add_argument('--debug-instructions', action='store_true', help='print every instruction; very slow')
    ap.add_argument('--verbose-commands', action='store_true', help='print every queued command and payload length')
    args = ap.parse_args()

    script_dir = Path(__file__).resolve().parent
    xdvd = Path(args.xdvd_json)
    if args.command_json:
        commands = commands_from_json(Path(args.command_json))
        source = Path(args.command_json)
        source_desc = f'command JSON: {source}'
    else:
        log_path = Path(args.log) if args.log else xdvd.with_suffix('.log')
        if not log_path.exists():
            raise SystemExit(f'[-] unlock log not found: {log_path}; pass --log PATH or --command-json PATH')
        commands = commands_from_unlock_log(log_path)
        source = log_path
        source_desc = f'unlock log: {log_path}'

    print('[V131] Persistent multi-command firmware session baseline.', flush=True)
    print(f'[V131] XDVD sidecar: {args.xdvd_json}', flush=True)
    if args.iso:
        print(f'[V131] ISO metadata: {args.iso}', flush=True)
    print(f'[V131] Command source: {source_desc}', flush=True)
    print(f'[V131] Commands queued: {len(commands)}; d108_hw_data_phase={0 if args.no_d108_hw_data_phase else 1}; debug_trace={1 if args.debug_trace else 0}', flush=True)
    if args.verbose_commands:
        for i, c in enumerate(commands, 1):
            data = parse_hex_bytes(c.get('data_in_hex') if c.get('direction') == 'in' else c.get('data_out_hex'))
            print(f"[V131] {i:02d} {c.get('tag')} {c.get('direction')} cdb={c.get('cdb_hex')} bytes={len(data)}", flush=True)

    stem = args.stem
    trace_csv = f'{stem}_trace.csv'
    config_json = script_dir / f'{stem}_config.json'
    output_txt = script_dir / f'{stem}.output.txt'
    verify_md = f'{stem}_verify.md'

    cfg = build_config(args, commands, stem)
    config_json.write_text(json.dumps(cfg, indent=2), encoding='utf-8')
    print(f'[OK] wrote persistent config: {config_json.name}', flush=True)

    run_checked([sys.executable, 'mn103s_emu.py', config_json.name], cwd=script_dir, stdout_file=output_txt)
    print(f'[OK] emulator output: {output_txt.name}', flush=True)

    res = subprocess.run([sys.executable, 'analyze_gdr_persistent_session.py', trace_csv, '--config', config_json.name, '--out', verify_md], cwd=str(script_dir))
    if res.returncode == 0:
        print('[OK] v131 persistent session completed exactly.', flush=True)
    else:
        print(f'[WARN] v131 verifier returned {res.returncode}; session incomplete or mismatched, but output/trace were written.', flush=True)
    print(f'[NEXT] Send {stem}.output.txt, {trace_csv}, {config_json.name}, and {verify_md}.', flush=True)
    return 0


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