"""Extract bounded, typed CW dictionaries from a reference MemoryPack cache.

Other cache schemas need not match. Each candidate must decode every row with
matching dictionary/row IDs, unique keys and bounded arrays. Require catalog
coverage for nodes and reject ambiguous largest candidates. This exports
reference data, not a claim of current Client value compatibility.
"""
import argparse
import gzip
import hashlib
import json
from pathlib import Path
import re
import struct


class Reader:
    def __init__(self, data, offset):
        self.data = data
        self.offset = offset

    def uint(self):
        value = struct.unpack_from('<I', self.data, self.offset)[0]
        self.offset += 4
        return value

    def byte(self):
        value = self.data[self.offset]
        self.offset += 1
        return value

    def array(self, fixed=False):
        count = self.uint()
        if count > 4096:
            raise ValueError('array exceeds extraction bound')
        values = []
        for _ in range(count):
            if fixed and self.byte() != 1:
                raise ValueError('fixed-value schema differs')
            values.append(self.uint())
        return values


def row(reader, kind):
    headers = {'portals': 5, 'node_templates': 7, 'penalty_rules': 5,
               'basic_bonuses': 5, 'combination_bonuses': 3}
    if reader.byte() != headers[kind]:
        raise ValueError('row schema differs')
    identity = reader.uint()
    result = {'ID': identity}
    if kind == 'portals':
        flag = reader.byte()
        if flag > 1:
            raise ValueError('invalid boolean')
        result.update(IfInBook=bool(flag), EffectParamList=reader.array(True),
                      ShowBonusIDList=reader.array(), PortalGameRefTrait=reader.array())
    elif kind == 'node_templates':
        result.update(StageID=reader.uint(), NodeType=reader.uint(),
                      ParamList=reader.array(), PenaltyBonusRuleID=reader.uint(),
                      IsAugment=reader.uint(), BasicGoldRewardNum=reader.uint())
        if result['NodeType'] not in range(6):
            raise ValueError('unknown node enum')
    elif kind == 'penalty_rules':
        result.update(ThresholdPosition=reader.uint(),
                      ThresholdFailPlayerHPPenalty=reader.uint(),
                      ThresholdPassBasicPlayerHPPenalty=reader.uint(),
                      ProgressPenaltyCoefficient=reader.uint())
        if result['ThresholdPosition'] > 100 or any(
                value > 1000000 for key, value in result.items() if key != 'ID'):
            raise ValueError('penalty values exceed known schema range')
    elif kind == 'basic_bonuses':
        result.update(Value=reader.uint(), BonusTypeParam=reader.uint(),
                      BonusTypeParamList=reader.array(), BonusType=reader.uint())
        # Preserve numeric bonus kinds added after the public source snapshot.
        # Consumers must reject kinds they do not implement, not drop rows.
        if result['BonusType'] > 255:
            raise ValueError('bonus kind exceeds observed schema bound')
    else:
        result.update(CombinationBonusList=reader.array(), BonusNumberList=reader.array())
        if len(result['CombinationBonusList']) != len(result['BonusNumberList']):
            raise ValueError('combination quantities differ')
    return result


def extract(data, kind, required):
    header = {'portals': 5, 'node_templates': 7, 'penalty_rules': 5,
              'basic_bonuses': 5, 'combination_bonuses': 3}[kind]
    pattern = re.compile(rb'(?P<key>.{4})' + bytes([header]) + rb'(?P=key)', re.S)
    candidates = []
    attempts = 0
    for match in pattern.finditer(data):
        start = match.start() - 4
        if start < 0:
            continue
        count = struct.unpack_from('<i', data, start)[0]
        if not 2 <= count <= 10000:
            continue
        attempts += 1
        if attempts > 50000:
            raise ValueError(f'{kind}: candidate bound exceeded')
        reader = Reader(data, start + 4)
        rows = {}
        try:
            for _ in range(count):
                key = reader.uint()
                value = row(reader, kind)
                if key != value['ID'] or key in rows:
                    raise ValueError('dictionary identity differs')
                rows[key] = value
                if reader.offset - start > 10000000:
                    raise ValueError('dictionary exceeds extraction bound')
            if required.issubset(rows):
                candidates.append((count, start, reader.offset, rows))
        except (ValueError, IndexError, struct.error):
            continue
    if not candidates:
        raise ValueError(f'{kind}: no complete matching dictionary')
    largest = max(candidate[0] for candidate in candidates)
    winners = [candidate for candidate in candidates if candidate[0] == largest]
    if len(winners) != 1:
        raise ValueError(f'{kind}: ambiguous full dictionary')
    count, start, end, rows = winners[0]
    print(f'{kind}: {count} rows, decompressed bytes {start}..{end}')
    return {'offset': start, 'length': end-start, 'rows': rows}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--resource', required=True, type=Path)
    parser.add_argument('--catalog', required=True, type=Path)
    parser.add_argument('--output', required=True, type=Path)
    args = parser.parse_args()
    packed = args.resource.read_bytes()
    data = gzip.decompress(packed)
    if len(data) > 256 * 1024 * 1024:
        raise ValueError('decompressed resource exceeds extraction bound')
    catalog = json.loads(args.catalog.read_text(encoding='utf-8-sig'))
    node_keys = {node['NodeTemplateID'] for node in catalog['stage_routes']}
    tables = {}
    for kind, required in [('node_templates', node_keys), ('portals', set()),
                           ('basic_bonuses', set()), ('combination_bonuses', set())]:
        tables[kind] = extract(data, kind, required)
    penalties = {value['PenaltyBonusRuleID']
                 for value in tables['node_templates']['rows'].values()
                 if value['PenaltyBonusRuleID']}
    tables['penalty_rules'] = extract(data, 'penalty_rules', penalties)
    output = {'format_version': 1, 'source_sha256': hashlib.sha256(packed).hexdigest(),
              'target_assembly_sha256': catalog['provenance']['game_assembly_sha256'],
              'compatibility': 'Reference dictionaries; target node ID coverage checked, values require Client comparison.',
              'tables': tables}
    args.output.write_text(json.dumps(output, indent=2)+'\n', encoding='utf-8')


if __name__ == '__main__':
    main()