# © 2022-2026 Richard Dymond (rjdymond@gmail.com)
#
# This file is part of SkoolKit.
#
# SkoolKit is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# SkoolKit is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# SkoolKit. If not, see <http://www.gnu.org/licenses/>.

import array

from skoolkit import SkoolKitError, write, write_line
from skoolkit.basic import TextReader
from skoolkit.pagingtracer import PagingTracer
from skoolkit.simulator import R1
from skoolkit.simutils import A, D, E, F, H, IXh, IXl, SP, PC, T, IFF
from skoolkit.traceutils import Registers, disassemble

DEC = tuple(tuple((
        v % 256,
        (v & 0xA8)                 # S.5.3...
        + (v == 0) * 0x40          # .Z......
        + (v % 16 == 0x0F) * 0x10  # ...H....
        + (v == 0x7F) * 0x04       # .....P..
        + 0x02                     # ......N.
        + c                        # .......C
    ) for v in range(-1, 255)
    ) for c in (0, 1)
)

DEC0 = DEC[0]

INC0 = tuple((
        v % 256,
        (v & 0xA8)                 # S.5.3.N.
        + (v == 256) * 0x40        # .Z......
        + (v % 16 == 0x00) * 0x10  # ...H....
        + (v == 0x80) * 0x04       # .....P..
    ) for v in range(1, 257)
)

class DataBlock:
    def __init__(self, data, start, end, keys, fast_load=True):
        self.data = data
        self.start = start
        self.end = end
        self.keys = keys
        self.fast_load = fast_load

    def adjust(self, max_index):
        self.start = min(self.start, max_index)
        self.end = min(self.end, max_index)

def _check_polarity(timings, polarity, edges, tstates, analysis):
    if timings.polarity is not None:
        ear = (len(edges) - 1) % 2
        if ear != timings.polarity ^ (polarity % 2):
            if analysis:
                analysis.append(f'{tstates:>10}  {ear:>3}  Polarity adjustment ({ear}->{ear ^ 1})')
            edges.append(tstates)

def get_edges(blocks, first_edge, polarity, analyse=False):
    edges = [first_edge]
    if polarity % 2:
        edges.append(first_edge)
    data_blocks = []
    tstates = first_edge
    tail = first_edge - 1
    keys = None
    if analyse:
        analysis = ['T-states    EAR  Description']
    else:
        analysis = ()
    byte_timings = {}

    for i, block in enumerate(blocks):
        timings = block.timings
        data = block.data
        keys = block.keys or keys

        # Pulses
        if timings.pulses:
            _check_polarity(timings, polarity, edges, tstates, analysis)
            for count, duration in timings.pulses:
                if analyse:
                    ear = (len(edges) - 1) % 2
                    if count == 1:
                        analysis.append(f'{tstates:>10}  {ear:>3}  Pulse ({duration} T-states)')
                    else:
                        analysis.append(f'{tstates:>10}  {ear:>3}  Tone ({count} x {duration} T-states)')
                for n in range(count):
                    tstates += duration
                    edges.append(tstates)

        # Data
        if data:
            _check_polarity(timings, polarity, edges, tstates, analysis)
            if analyse:
                if timings.used_bits < 8:
                    bits = f' + {timings.used_bits} bits'
                    data_len = len(data) - 1
                else:
                    bits = ''
                    data_len = len(data)
                ear = (len(edges) - 1) % 2
                zero = ','.join(str(d) for d in timings.zero)
                one = ','.join(str(d) for d in timings.one)
                analysis.append(f'{tstates:>10}  {ear:>3}  Data ({data_len} bytes{bits}; {zero}/{one} T-states)')
            start = len(edges) - 1

            if 0 in timings.zero or 0 in timings.one:
                p = q = 0
                for k, b in enumerate(data, 1):
                    for j in range(8 if k < len(data) else timings.used_bits):
                        for d in timings.one if b & 0x80 else timings.zero:
                            if d:
                                tstates += d
                                if p == q:
                                    edges.append(tstates)
                                    q = 1 - q
                                else:
                                    edges[-1] += d
                            p = 1 - p
                        b *= 2
            else:
                key = (timings.zero, timings.one)
                b_timings = byte_timings.get(key)
                if b_timings is None:
                    b_timings = [None] * 256
                    for v in range(256):
                        offsets = []
                        b = v
                        for j in range(8):
                            offsets.extend(timings.one if b & 0x80 else timings.zero)
                            b *= 2
                        b_timings[v] = offsets
                    byte_timings[key] = b_timings
                for b in data[:-1]:
                    for d in b_timings[b]:
                        tstates += d
                        edges.append(tstates)
                bt = b_timings[data[-1]]
                for d in bt[:(len(bt) * timings.used_bits) // 8]:
                    tstates += d
                    edges.append(tstates)

            # Tail pulse
            if timings.tail:
                if analyse:
                    ear = (len(edges) - 1) % 2
                    analysis.append(f'{tstates:>10}  {ear:>3}  Tail pulse ({timings.tail} T-states)')
                tstates += timings.tail
                tail = tstates
                edges.append(tstates)

            if 0 in timings.zero or 0 in timings.one:
                # This is sample data as opposed to byte values, so ensure it's
                # not fast-loaded
                data_blocks.append(DataBlock(data, len(edges) - 1, len(edges) - 1, keys, False))
                keys = None
            else:
                data_blocks.append(DataBlock(data, start, len(edges) - 1, keys))
                keys = None
        elif timings.data or (i == len(blocks) - 1 and timings.pulses):
            # If this block contains a Direct Recording, or is the last block
            # on the tape and contains pulses but no data, add a data block to
            # ensure that the pulses are read
            data_blocks.append(DataBlock((), len(edges) - 1, len(edges) - 1, keys, False))
            keys = None

        # Pause
        if i + 1 < len(blocks) and timings.pause:
            _check_polarity(timings, polarity, edges, tstates, analysis)
            if analyse:
                ear = (len(edges) - 1) % 2
                analysis.append(f'{tstates:>10}  {ear:>3}  Pause ({timings.pause} T-states)')
            tstates += timings.pause

    if edges[-1] == tail:
        edges.pop()
        if analysis:
            analysis.pop()
        if data_blocks:
            data_blocks[-1].adjust(len(edges) - 1)

    for line in analysis:
        print(line)

    return edges, data_blocks

class LoadTracer(PagingTracer):
    def __init__(self, simulator, blocks, config, draw):
        self.tsl_misses = 0
        self.dec_a_jr_hits = 0
        self.dec_a_jp_hits = 0
        self.dec_a_misses = 0
        self.simulator = simulator
        self.frame_duration = simulator.frame_duration
        self.edges, self.blocks = get_edges(blocks, config['first_edge'], config['polarity'])
        self.keys = None
        self.pause = config['pause']
        self.in_min_addr = config['in_min_addr']
        self.accelerators = config['accelerators']
        self.accel_dec_a = config['accelerate_dec_a']
        if hasattr(simulator, 'opcodes') and self.accel_dec_a: # pragma: C no cover
            dec_a_jr = self.accel_dec_a & 1
            dec_a_jp = self.accel_dec_a & 2
            simulator.opcodes[0x3D] = self.dec_a(dec_a_jr, dec_a_jp)
        self.list_accelerators = config['list_accelerators']
        self.block_index = 0
        self.block_data_index = self.blocks[0].start
        self.max_index = len(self.edges) - 1
        self.stop = config['stop']
        self.flash_load = config['fast_load']
        self.finish_tape = config['finish_tape']
        self.timeout = config['timeout']
        self.tracefile = config['tracefile']
        self.trace_line = config['trace_line']
        self.prefix = config['prefix']
        self.byte_fmt = config['byte_fmt']
        self.word_fmt = config['word_fmt']
        self.draw = draw
        if draw:
            self.write_port = self.write_port_with_border_list
        self.text = TextReader()
        self.state = [
            0,                  # state[0]: next edge (timestamp)
            0,                  # state[1]: edge index
            0,                  # state[2]: end of tape reached
            self.blocks[0].end, # state[3]: index of final edge in current block
            0,                  # state[4]: tape running
            0,                  # state[5]: custom loader detected
            0,                  # state[6]: tape end time
            1,                  # state[7]: data block not yet announced
            0,                  # state[8]: time of next interrupt
            0,                  # state[9]: last frame drawn
        ]
        if hasattr(simulator, 'load'): # pragma: Python no cover
            self.edges = array.array('Q', self.edges)
            self.state = array.array('Q', self.state)
        self.read_port = self._read_port()

    def run(self, border, out7ffd, outfffd, ay, outfe):
        is128k = len(self.simulator.memory) == 0x20000
        if is128k:
            self.out7ffd = out7ffd # 128K ROM 0/1
        else:
            self.out7ffd = 0x10 # Signal: 48K ROM always
        self.outfffd = outfffd
        self.ay = ay
        self.outfe = outfe
        self.keys = None
        simulator = self.simulator
        memory = simulator.memory
        registers = simulator.registers
        tracefile = self.tracefile
        if tracefile:
            trace_line = self.trace_line
            prefix = self.prefix
            byte_fmt = self.byte_fmt
            word_fmt = self.word_fmt
            r = Registers(registers)
            m = memory
        draw = self.draw
        if draw and isinstance(border, int):
            self.border = [(0, border)]
        else:
            self.border = border

        if hasattr(simulator, 'load'): # pragma: Python no cover
            if tracefile:
                df = lambda pc: disassemble(memory, pc, prefix, byte_fmt, word_fmt)[0]
                tf = lambda pc, i, t0: tracefile.write(trace_line.format(pc=pc, i=i, r=r, t=t0, m=m))
            else:
                df = tf = None
            stop_cond = simulator.load(self.stop, self.flash_load, self.finish_tape, self.timeout, draw, df, tf)
            pc = registers[24]
        else: # pragma: C no cover
            opcodes = simulator.opcodes
            frame_duration = self.frame_duration
            int_active = simulator.int_active
            pc = registers[24]
            progress = 0
            edges = self.edges
            tape_length = edges[-1] // 1000
            max_index = self.max_index
            tstates = registers[25]
            state = self.state
            stop = self.stop
            fast_load = self.flash_load
            finish_tape = self.finish_tape
            timeout = self.timeout
            state[8] = ((tstates + frame_duration - int_active) // frame_duration) * frame_duration
            while True:
                t0 = tstates
                if tracefile:
                    i = disassemble(memory, pc, prefix, byte_fmt, word_fmt)[0]
                    opcodes[memory[pc]]()
                    tracefile.write(trace_line.format(pc=pc, i=i, r=r, t=t0, m=m))
                else:
                    opcodes[memory[pc]]()
                tstates = registers[25]

                if tstates >= state[8]:
                    if tstates < state[8] + int_active:
                        if registers[26]:
                            simulator.accept_interrupt(registers, memory, pc)
                            tstates = registers[25]
                        if draw:
                            frame = tstates // frame_duration
                            if frame > state[9]:
                                if is128k:
                                    scr = memory.memory[1][:6912]
                                else:
                                    scr = memory[16384:23296]
                                if not draw(scr, frame, self.border):
                                    stop_cond = 6
                                    break
                                state[9] = frame
                    else:
                        state[8] = ((tstates + frame_duration - int_active) // frame_duration) * frame_duration
                        state[9] = tstates // frame_duration

                if state[4] and tstates >= state[0]:
                    index = state[1]
                    while index < max_index and edges[index + 1] < tstates:
                        index += 1
                    state[1] = index
                    if index == max_index:
                        # Allow 1ms for the final edge on the tape to be read
                        if tstates - edges[index] > 3500:
                            self.stop_tape(tstates)
                    elif index > state[3]:
                        # Pause tape between blocks
                        self.next_block(tstates)
                        if self.keys:
                            stop_cond = 5
                            break
                    else:
                        state[0] = edges[index + 1]
                        if index < state[3]:
                            p = edges[index] // tape_length
                            if p > progress:
                                write(f'[{p/10:5.1f}%]\x08\x08\x08\x08\x08\x08\x08\x08')
                                progress = p

                pc = registers[24]

                if pc == stop and (state[2] or not finish_tape):
                    stop_cond = 0
                    break

                if pc == 0x0556 and self.out7ffd & 0x10 and fast_load and self.fast_load(simulator):
                    state[1] = state[3]
                    if state[1] == max_index:
                        # Final block, so stop the tape
                        self.stop_tape(tstates)
                    else:
                        # Otherwise continue to play the tape until this block's
                        # 'pause' period (if any) has elapsed
                        state[4] = 1 # Signal: tape is running
                        registers[25] = state[0] = edges[state[1]]
                        state[8] = ((registers[25] + frame_duration - int_active) // frame_duration) * frame_duration
                        state[9] = 0
                    pc = registers[24]
                    tstates = registers[25]
                else:
                    if state[2] and stop is None:
                        # The tape has ended and no stop address is set
                        if state[5]:
                            # Custom loader was detected
                            stop_cond = 1
                            break
                        if pc > 0x3FFF:
                            # PC in RAM
                            stop_cond = 2
                            break
                        if tstates - self.state[6] > 3500000:
                            # Tape ended 1 second ago
                            stop_cond = 3
                            break
                    if tstates > timeout:
                        stop_cond = 4
                        break

        if stop_cond == 0:
            write_line(f'Simulation stopped (PC at start address): PC={pc}')
        elif stop_cond == 1:
            write_line(f'Simulation stopped (end of tape): PC={pc}')
        elif stop_cond == 2:
            write_line(f'Simulation stopped (PC in RAM): PC={pc}')
        elif stop_cond == 3:
            write_line(f'Simulation stopped (tape ended 1 second ago): PC={pc}')
        elif stop_cond == 4:
            write_line(f'Simulation stopped (timed out): PC={pc}')
        elif stop_cond == 5:
            write_line('Tape paused')
        elif stop_cond == 6:
            write_line('Simulation stopped (screen closed)')

    def dec_a(self, dec_a_jr, dec_a_jp): # pragma: C no cover
        # Speed up any 'DEC A: JR/JP NZ,$-1' loop if configured to do so, and
        # also count hits and misses
        registers = self.simulator.registers
        memory = self.simulator.memory
        def func():
            pc = registers[24]
            if registers[26] == 0:
                if dec_a_jr and memory[(pc + 1) % 65536] == 0x20 and memory[(pc + 2) % 65536] == 0xFD:
                    self.dec_a_jr_hits += 1
                    a = registers[0]
                    if a == 0:
                        a = 256
                    registers[0] = 0
                    registers[1] = 0x42 + (registers[1] % 2)
                    r = registers[15]
                    registers[15] = (r & 0x80) + ((r + a * 2) % 128)
                    registers[25] += 16 * a - 5
                    registers[24] = (pc + 3) % 65536
                    return
                if dec_a_jp and memory[(pc + 1) % 65536] == 0xC2 and memory[(pc + 2) % 65536] == pc % 256 and memory[(pc + 3) % 65536] == pc // 256:
                    self.dec_a_jp_hits += 1
                    a = registers[0]
                    if a == 0:
                        a = 256
                    registers[0] = 0
                    registers[1] = 0x42 + (registers[1] % 2)
                    r = registers[15]
                    registers[15] = (r & 0x80) + ((r + a * 2) % 128)
                    registers[25] += 14 * a
                    registers[24] = (pc + 4) % 65536
                    return
                self.dec_a_misses += 1
            registers[0], registers[1] = DEC[registers[1] % 2][registers[0]]
            registers[15] = R1[registers[15]]
            registers[25] += 4
            registers[24] = (pc + 1) % 65536
        return func

    def _read_port(self): # pragma: C no cover
        in_min_addr = self.in_min_addr
        state = self.state
        edges = self.edges
        blocks = self.blocks
        max_index = self.max_index
        accelerators = list(self.accelerators)
        memory = self.simulator.memory
        def func(registers, port):
            if port & 0xFF == 0xFE:
                pc = registers[24]
                if pc >= in_min_addr or (0x0562 <= pc <= 0x05F1 and self.out7ffd & 0x10):
                    state[5] = 1 # Signal: custom loader detected
                    index = state[1]
                    if state[7] and not state[2]:
                        state[7] = 0 # Signal: data block announced
                        state[4] = 1 # Signal: tape is running
                        registers[T] = edges[index]
                        state[8] = ((registers[T] + self.frame_duration - self.simulator.int_active) // self.frame_duration) * self.frame_duration
                        state[9] = 0
                        length = len(blocks[self.block_index].data)
                        if length:
                            write_line(f'Data ({length} bytes)')
                    elif index == max_index:
                        # Final edge, so stop the tape
                        self.stop_tape(registers[T])
                    elif state[4] and registers[26] == 0 and index < state[3] - 1:
                        loops = 0
                        for i, acc in enumerate(accelerators):
                            if memory[pc - acc.c0:pc + acc.c1] == acc.code:
                                acc.hits += 1
                                if acc.ear_mask:
                                    ffwd = registers[acc.ear] & acc.ear_mask == ((index - acc.polarity) % 2) * acc.ear_mask
                                else:
                                    # Polarity-sensitive loop
                                    ffwd = (index - acc.polarity) % 2
                                if ffwd and state[0] > registers[25]:
                                    counter = registers[acc.counter]
                                    if acc.inc:
                                        # INC r
                                        loops = min((state[0] - registers[25]) // acc.loop_time + 1, 255 - counter)
                                    else:
                                        # DEC r
                                        loops = min((state[0] - registers[25]) // acc.loop_time + 1, counter - 1)
                                    if loops:
                                        if acc.inc:
                                            # INC r
                                            registers[acc.counter], registers[1] = INC0[counter + loops - 1]
                                        else:
                                            # DEC r
                                            registers[acc.counter], registers[1] = DEC0[counter - loops + 1]
                                        r = registers[15]
                                        registers[15] = (r & 0x80) + ((r + acc.loop_r_inc * loops) % 0x80)
                                        registers[25] += acc.loop_time * loops
                                        if registers[25] > state[0]:
                                            index += 1
                                if i:
                                    # Move the selected accelerator to the beginning of the
                                    # list so that it can be found quicker next time
                                    accelerators.remove(acc)
                                    accelerators.insert(0, acc)
                                break
                        else:
                            self.tsl_misses += 1
                    if index % 2 == 0:
                        return 191
            elif port & 0xC002 == 0xC000 and self.outfffd < 16:
                ay_reg = self.outfffd
                if ay_reg == 14 and registers[24] == 0x08B2:
                    # Avoid an infinite loop at 0x08AF in the 128K ROM:
                    #   $08AF CALL $05D6  ; Check for BREAK.
                    #   $08B2 IN A,(C)    ; Read AY register 14 (BC=$FFFD).
                    #   $08B4 AND $40     ; Ready to send data?
                    #   $08B6 JR NZ,$08AF ; Jump back if not (bit 6 set).
                    return 0
                return self.ay[ay_reg]
            return 255
        return func

    def next_block(self, tstates):
        self.block_index += 1
        if self.block_index >= len(self.blocks):
            self.stop_tape(tstates)
        else:
            self.state[1] = self.state[3] + 1
            self.state[0] = self.edges[self.state[1]]
            block = self.blocks[self.block_index]
            self.block_data_index = block.start
            self.keys = block.keys
            self.state[3] = self.blocks[self.block_index].end
            self.state[4] = int(not self.pause) # Pause tape unless configured not to
            self.state[7] = 1 # Signal: data block not yet announced

    def stop_tape(self, tstates):
        self.block_index = len(self.blocks)
        self.state[2] += 1 # Signal: end of tape reached
        if self.state[2] == 1:
            write_line('Tape finished')
            self.state[6] = tstates # Set tape end time
        self.state[4] = 0 # Signal: tape is not running

    def fast_load(self, simulator):
        registers = simulator.registers
        while self.block_data_index <= self.state[1] < self.max_index:
            self.next_block(registers[T])
        if self.block_index < len(self.blocks):
            data_block = self.blocks[self.block_index]
        else:
            raise SkoolKitError("Failed to fast load block: unexpected end of tape")
        if not data_block.fast_load:
            return False

        memory = simulator.memory
        ix = registers[IXl] + 256 * registers[IXh] # Start address
        de = registers[E] + 256 * registers[D] # Block length
        a = registers[A]
        block = data_block.data
        data_len = len(block) - 2

        # Exchange AF register pairs (as done at 0x0557), disable interrupts
        # (as done at 0x0559), and preload the machine stack with 0x053F (as
        # done at 0x055E)
        registers[0:2], registers[16:18] = registers[16:18], registers[0:2]
        registers[IFF] = 0
        sp = (registers[SP] - 2) % 65536
        registers[SP] = sp
        if sp > 0x3FFF:
            memory[sp] = 0x3F
        sp = (sp + 1) % 65536
        if sp > 0x3FFF:
            memory[sp] = 0x05

        if a == block[0]:
            skipped = ''
        else:
            skipped = ' [skipped]'
        if block[0] == 0 and data_len >= 17 and block[1] <= 3:
            name = self.text.get_text(block[2:12])
            if block[1] == 3:
                write_line(f'Bytes: {name}{skipped}')
            elif block[1] == 2:
                write_line(f'Character array: {name}{skipped}')
            elif block[1] == 1:
                write_line(f'Number array: {name}{skipped}')
            elif block[1] == 0:
                write_line(f'Program: {name}{skipped}')
        elif skipped:
            write_line(f'Data block ({data_len} bytes){skipped}')
        else:
            write_line(f'Fast loading data block: {ix},{de}')

        if skipped:
            registers[F] = 0x00 # Reset ZF, reset CF: flag byte mismatch
            registers[H] = 0
        else:
            addr = ix
            if de <= data_len:
                length = de
                ix += de
                de = 0
                check_parity = True
            else:
                length = data_len + 1
                ix += data_len + 1
                de -= data_len + 1
                check_parity = False
                registers[F] = 0x40 # Set ZF, reset CF: edge detection failed
            i = 1
            parity = block[0]
            while length:
                if addr > 0x3FFF:
                    memory[addr] = block[i]
                parity ^= block[i]
                addr = (addr + 1) % 65536
                i += 1
                length -= 1
            if check_parity:
                a = parity ^ block[i]
                registers[A] = a
                registers[F] = (a == 1) * 0x40 + (a == 0)
            registers[IXh] = (ix >> 8) & 0xFF
            registers[IXl] = ix & 0xFF
            registers[D] = (de >> 8) & 0xFF
            registers[E] = de & 0xFF

        registers[PC] = 0x05E2
        self.state[7] = 0 # Signal: data block announced
        return True
