#!/usr/bin/env python3
"""
sceneryRender - render an FS2 scenery snapshot from the extracted database.

Walks the chunk5 bytecode using the polygon-rendering subset of opcodes
($07 EnterLocalFrame, $24 PushOriginWithStash, $12 SetColor, $40/$41
xform-B emit, $01/$02 xform-A emit, $31/$42 cache fill, $32/$33/$35
cached vertex). Uses MAME's captured matrix for the camera transform so
the output is bit-equivalent to running the original on real hardware.

Usage:
  sceneryRender.py FS2.1.json [--out path.ppm]
"""

import argparse
import json
import os
import struct
import sys

# Viewport (Apple II hires above the panel).
NATIVE_WIDTH = 280
NATIVE_HEIGHT = 192
VIEWPORT_BOTTOM = 99       # Last row of the scenery viewport.

# Camera state captured from MAME at the Meigs boot moment.
# $5C/$5D = camera X (scenery units), $64/$65 = camera Z, $5E/$5F = altitude.
CAMERA_X_SCENERY = 287
CAMERA_Y_SCENERY = 25       # = altitude in metres for MAME boot
CAMERA_Z_SCENERY = 804

# 3x3 view matrix at boot. Captured from MAME runtime ZP $78-$89.
# Q1.14 entries (signed 16-bit, bit 14 = 1.0).
MATRIX = (
    (16382,    0,    0),
    (    0, 32760,  100),
    (    0,  -401, 8190),
)

# Palette codes -> 24-bit RGB (chunk5 ToHiresColorTable mapping).
COLORS = {
    0x00: (0,   0,   0),
    0x01: (63,  125, 47),    # ground-day green
    0x02: (42,  93,  168),   # water blue
    0x03: (63,  125, 47),
    0x04: (42,  93,  168),
    0x05: (0,   0,   0),
    0x06: (139, 134, 128),   # runway grey
    0x07: (160, 96,  64),    # building tan/red
    0x08: (0,   0,   0),
    0x09: (224, 224, 224),
    0x0A: (156, 122, 79),
    0x0B: (42,  85,  32),
    0x0C: (110, 90,  63),
    0x0D: (240, 240, 240),   # white
    0x0E: (200, 224, 240),
    0x0F: (176, 168, 120),   # CITY tan
}
SKY_COLOR    = (107, 163, 220)
GROUND_COLOR = (63,  125, 47)
DEFAULT_LINE_COLOR = (240, 240, 240)


# -- bytecode walker that renders ---------------------------------------------

class State:
    def __init__(self):
        # Camera-vs-section delta ($66/$67=X, $68/$69=Y, $6A/$6B=Z).
        # Per port: at boot, $66=$6A=0, $68 = camera Y in scenery units.
        self.deltaX = 0
        self.deltaY = CAMERA_Y_SCENERY
        self.deltaZ = 0
        # Cached vertex pool ($0140 + idx*8). Each entry: (x, y, z, outcode).
        self.cache = [None] * 64
        # V1 / V2 holding cells.
        self.v1 = (0, 0, 0, 0)        # (x, y, z, outcode)
        self.v2 = (0, 0, 0, 0)
        # Color (palette code).
        self.colorCode = 0x0D         # white default
        # Day-only suppression (chunk5 $1B clears, $1C sets at night).
        self.dayOnlySkip = False
        # Polygon-state $D3 (= outcode AND of accumulated primary verts).
        self.d3 = 0
        # All emitted segments: (x1,y1,z1, x2,y2,z2, color_code).
        self.segments = []
        # All cached-vertex emits (for stats).
        self.cacheFills = 0


# -- transform & projection ---------------------------------------------------

def transform7EBC(state, vx, vz):
    """xform-B: subtract camera delta then multiply by view matrix.

    chunk5 TransformVertex7EBC reads X and Z from the stream (no Y; Y comes
    from the section base, accumulated by EnterLocalFrame). The result is
    the camera-relative (X, Y, Z) ready for projection.
    """
    dx = vx - state.deltaX
    dy = -state.deltaY               # local Y = 0 - camera Y
    dz = vz - state.deltaZ
    return _matmul(dx, dy, dz)


def transform80C5(state, vx, vy, vz):
    """xform-A: 6-byte X/Y/Z stream; Y is per-vertex."""
    dx = vx - state.deltaX
    dy = vy - state.deltaY
    dz = vz - state.deltaZ
    return _matmul(dx, dy, dz)


def _matmul(dx, dy, dz):
    # Q1.14 matrix * Q0 vector. Output Q0 (world units).
    # chunk5 normalises by >> 14 (or 16 for some paths); use >> 14 to match.
    rx = (MATRIX[0][0] * dx + MATRIX[0][1] * dy + MATRIX[0][2] * dz) >> 14
    ry = (MATRIX[1][0] * dx + MATRIX[1][1] * dy + MATRIX[1][2] * dz) >> 14
    rz = (MATRIX[2][0] * dx + MATRIX[2][1] * dy + MATRIX[2][2] * dz) >> 14
    return (rx, ry, rz)


def classifyOutcode(v):
    """6-bit frustum outcode (bit 7 = behind-camera valid flag)."""
    x, y, z = v
    code = 0
    if z <= 0:
        code |= 0x20
        return code | 0x80
    if x < -z: code |= 0x01
    if x >  z: code |= 0x02
    if y < -z: code |= 0x04
    if y >  z: code |= 0x08
    return code | 0x80


def project(v):
    x, y, z = v
    if z <= 0:
        return None
    qx = (x * 128) // z
    qy = (y * 128) // z
    cx = NATIVE_WIDTH    // 2
    cy = VIEWPORT_BOTTOM // 2
    sx = cx + (qx * (NATIVE_WIDTH    // 2)) // 0x7F
    sy = cy - (qy * (VIEWPORT_BOTTOM // 2)) // 0x7F
    return (sx, sy)


def clipNearPlane(a, b):
    if a[2] > 0 and b[2] > 0:
        return a, b
    if a[2] <= 0 and b[2] <= 0:
        return None
    behind, front = (a, b) if a[2] <= 0 else (b, a)
    denom = front[2] - behind[2]
    if denom <= 0:
        return None
    t = ((1 - behind[2]) << 8) // denom
    if t < 0: t = 0
    if t > 256: t = 256
    bx = behind[0] + ((front[0] - behind[0]) * t) // 256
    by = behind[1] + ((front[1] - behind[1]) * t) // 256
    return (a, (bx, by, 1)) if b is behind else ((bx, by, 1), b)


# -- bytecode interpreter ----------------------------------------------------

def runSection(secOps, state):
    """Walk a section's already-decoded ops in pc order. The extractor's
    walker followed both branches of every conditional, so we just iterate
    in pc order and process the rendering-relevant ops.
    """
    for op in sorted(secOps, key=lambda o: o['pc']):
        code = op['op']
        if code == 0x12:                     # SetColor
            state.colorCode = op['color_code']
        elif code == 0x1B:                   # ModeWhite
            state.dayOnlySkip = False
            state.colorCode = 0x0D
        elif code == 0x1C:                   # DayOnly
            pass    # daytime render: don't suppress
        elif code == 0x07:                   # EnterLocalFrame
            # Layout: $07, base, x*3, y*3, z*3, dx*2, dy*2 = 14 bytes.
            # For now: extract section base and shift camera delta.
            b = op['bytes']
            if len(b) >= 8:
                bx = _signed16(b[2], b[3])
                by = _signed16(b[4], b[5])
                bz = _signed16(b[6], b[7])
                state.deltaX += bx
                state.deltaY += by
                state.deltaZ += bz
        elif code == 0x24:                   # PushOriginWithStash
            # Layout: $24, dx*2, dy*2, dz*2 = 8 bytes (approx).
            # Pushes a new origin onto the section stack.
            b = op['bytes']
            if len(b) >= 8:
                bx = _signed16(b[2], b[3])
                by = _signed16(b[4], b[5])
                bz = _signed16(b[6], b[7])
                state.deltaX += bx
                state.deltaY += by
                state.deltaZ += bz
        elif code == 0x40:                   # EmitV1Xform7EBC silent
            if 'vx' not in op or 'vz' not in op: continue
            v = transform7EBC(state, op['vx'], op['vz'])
            state.v1 = (*v, classifyOutcode(v))
        elif code == 0x41:                   # EmitV2Xform7EBC + draw
            if 'vx' not in op or 'vz' not in op: continue
            v = transform7EBC(state, op['vx'], op['vz'])
            state.v2 = (*v, classifyOutcode(v))
            _emitSegment(state)
            state.v1 = state.v2
        elif code == 0x01:                   # EmitV1Xform80C5 silent
            if not all(k in op for k in ('vx','vy','vz')): continue
            v = transform80C5(state, op['vx'], op['vy'], op['vz'])
            state.v1 = (*v, classifyOutcode(v))
        elif code == 0x02:                   # EmitV2Xform80C5 + draw
            if not all(k in op for k in ('vx','vy','vz')): continue
            v = transform80C5(state, op['vx'], op['vy'], op['vz'])
            state.v2 = (*v, classifyOutcode(v))
            _emitSegment(state)
            state.v1 = state.v2
        elif code == 0x42:                   # cache fill xform-B
            if 'cache_idx' not in op or 'vx' not in op: continue
            v = transform7EBC(state, op['vx'], op['vz'])
            state.cache[op['cache_idx']] = (*v, classifyOutcode(v))
            state.cacheFills += 1
        elif code == 0x31:                   # cache fill xform-A
            if not all(k in op for k in ('cache_idx','vx','vy','vz')): continue
            v = transform80C5(state, op['vx'], op['vy'], op['vz'])
            state.cache[op['cache_idx']] = (*v, classifyOutcode(v))
            state.cacheFills += 1
        elif code == 0x32:                   # cached V1 (silent)
            if 'cache_idx' not in op: continue
            cv = state.cache[op['cache_idx']]
            if cv is not None:
                state.v1 = cv
        elif code == 0x33:                   # cached V2 + draw V1->V2
            if 'cache_idx' not in op: continue
            cv = state.cache[op['cache_idx']]
            if cv is not None:
                state.v2 = cv
                _emitSegment(state)
                state.v1 = state.v2
        elif code == 0x35:                   # cached + plot pixel
            if 'cache_idx' not in op: continue
            cv = state.cache[op['cache_idx']]
            if cv is not None:
                state.v1 = cv


def _emitSegment(state):
    if state.dayOnlySkip:
        return
    v1 = state.v1
    v2 = state.v2
    # Outcode AND -> entirely off-screen on one side.
    if (v1[3] & v2[3] & 0x3F) != 0:
        return
    state.segments.append((v1[0], v1[1], v1[2], v2[0], v2[1], v2[2], state.colorCode))


def _signed16(lo, hi):
    v = lo | (hi << 8)
    return v - 0x10000 if v >= 0x8000 else v


# -- framebuffer + draw ------------------------------------------------------

class FB:
    def __init__(self):
        self.pix = bytearray(NATIVE_WIDTH * NATIVE_HEIGHT * 3)

    def fillRow(self, y, color):
        for x in range(NATIVE_WIDTH):
            self.setPixel(x, y, color)

    def setPixel(self, x, y, color):
        if 0 <= x < NATIVE_WIDTH and 0 <= y < NATIVE_HEIGHT:
            i = (y * NATIVE_WIDTH + x) * 3
            self.pix[i] = color[0]
            self.pix[i+1] = color[1]
            self.pix[i+2] = color[2]

    def line(self, x1, y1, x2, y2, color):
        # Bresenham, clipped to viewport.
        dx = abs(x2 - x1); sx = 1 if x1 < x2 else -1
        dy = -abs(y2 - y1); sy = 1 if y1 < y2 else -1
        err = dx + dy
        steps = 0
        while True:
            self.setPixel(x1, y1, color)
            if x1 == x2 and y1 == y2: break
            e2 = 2 * err
            if e2 >= dy: err += dy; x1 += sx
            if e2 <= dx: err += dx; y1 += sy
            steps += 1
            if steps > 2000:        # safety cap
                break

    def writePpm(self, path):
        with open(path, 'wb') as f:
            f.write(b'P6\n%d %d\n255\n' % (NATIVE_WIDTH, NATIVE_HEIGHT))
            f.write(bytes(self.pix))


# -- main --------------------------------------------------------------------

def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("db", help="path to *.json (e.g. extracted_db/FS2.1.json)")
    p.add_argument("--out", default="port/screenshots/db_render.ppm")
    p.add_argument("--cam-x", type=int, default=CAMERA_X_SCENERY)
    p.add_argument("--cam-z", type=int, default=CAMERA_Z_SCENERY)
    p.add_argument("--cam-y", type=int, default=CAMERA_Y_SCENERY)
    p.add_argument("--all-sections", action="store_true",
                   help="walk every section, not just dispatcher-cull-passing")
    args = p.parse_args()

    db = json.load(open(args.db))
    print(f"Loaded {len(db['sections'])} sections from {db['source']}",
          file=sys.stderr)

    # Find dispatcher and decide which sections to render.
    disp = next(s for s in db['sections'] if s['sid'] == 0x2C)
    sectionsToRender = []
    if args.all_sections:
        sectionsToRender = list(db['sections'])
    else:
        ops = sorted(disp['ops'], key=lambda o: o['pc'])
        i = 0
        while i < len(ops):
            op = ops[i]
            if op['op'] == 0x13 and i + 1 < len(ops) and ops[i+1]['op'] == 0x0D:
                b = op['bytes']
                bound = b[3] | (b[4] << 8)
                refY = _signed16(b[5], b[6])
                refX = _signed16(b[7], b[8])
                if (abs(refX - args.cam_x) <= bound
                    and abs(refY - args.cam_z) <= bound):
                    sid = ops[i+1]['bytes'][1]
                    sec = next((s for s in db['sections'] if s['sid'] == sid), None)
                    if sec is not None:
                        sectionsToRender.append(sec)
                        print(f"  cull pass: load sid ${sid:02X}", file=sys.stderr)
                i += 2
                continue
            i += 1
        # Always include the dispatcher so any inline polygons render too.
        sectionsToRender.insert(0, disp)

    print(f"Rendering {len(sectionsToRender)} sections", file=sys.stderr)

    # Render.
    fb = FB()
    # Sky/ground fill at horizon row (= cy = VIEWPORT_BOTTOM/2).
    for y in range(NATIVE_HEIGHT):
        if y <= VIEWPORT_BOTTOM // 2:
            fb.fillRow(y, SKY_COLOR)
        elif y < VIEWPORT_BOTTOM:
            fb.fillRow(y, GROUND_COLOR)
        else:
            fb.fillRow(y, (0, 0, 0))    # panel area

    state = State()
    for sec in sectionsToRender:
        runSection(sec['ops'], state)

    print(f"Emitted {len(state.segments)} 3D segments", file=sys.stderr)

    # Project & draw.
    drawn = 0
    for v1x, v1y, v1z, v2x, v2y, v2z, colCode in state.segments:
        clip = clipNearPlane((v1x, v1y, v1z), (v2x, v2y, v2z))
        if clip is None: continue
        a, b = clip
        sa = project(a); sb = project(b)
        if sa is None or sb is None: continue
        col = COLORS.get(colCode & 0x0F, DEFAULT_LINE_COLOR)
        fb.line(sa[0], sa[1], sb[0], sb[1], col)
        drawn += 1

    print(f"Drew {drawn} 2D segments to {args.out}", file=sys.stderr)
    fb.writePpm(args.out)


if __name__ == "__main__":
    main()