"""Run in IDA via File > Script file; choose the generated Replay MAP file."""
from __future__ import annotations
import hashlib
import json
import re
from collections import Counter
from pathlib import Path

MAP_SYMBOL=re.compile(r'^\s*([0-9A-Fa-f]{4}):([0-9A-Fa-f]{8})\s+(\S+)\s+([0-9A-Fa-f]{16})\s')
GENERIC=re.compile(r'^(?:(?:j_)*(?:sub|nullsub|loc|locret|unk|off|byte|word|dword|qword|oword|xmmword|ymmword|zmmword|asc|stru|jpt|def|algn|flt|dbl|tbyte|packreal)_[0-9a-f]+(?:_[0-9]+)?|function(?:_[0-9]+)?)$',re.I)


def read_package(map_path):
    map_path=Path(map_path)
    manifest_path=map_path.with_suffix('.symbols.json')
    manifest=json.loads(manifest_path.read_text(encoding='utf-8'))
    if manifest.get('schema_version')!=1 or not manifest.get('symbols'):
        raise ValueError('Unsupported or empty symbol manifest')
    raw=map_path.read_bytes()
    if hashlib.sha256(raw).hexdigest()!=manifest['map_sha256']:
        raise ValueError('MAP SHA-256 does not match manifest')
    records=[]
    for line in raw.decode('utf-8').splitlines():
        match=MAP_SYMBOL.match(line)
        if match:
            section,offset,name,va=match.groups()
            records.append((int(section,16),int(offset,16),name,int(va,16)))
    wanted=[(s['section'],s['offset'],s['map_name'],s['va']) for s in manifest['symbols']]
    if records!=wanted:
        raise ValueError('MAP symbols do not match manifest addresses/names')
    sections={s['index']:s for s in manifest['sections']}
    for item in manifest['symbols']:
        section=sections[item['section']]
        if item['rva']!=section['rva']+item['offset'] or item['va']!=manifest['image_base']+item['rva'] or not 0<=item['rva']<manifest['image_size']:
            raise ValueError('Symbol section/RVA/VA mismatch')
    return manifest


def apply_package(map_path,adapter,replace_existing=False):
    """Adapter makes rebase/build/conflict behavior independently testable."""
    manifest=read_package(map_path)
    if adapter.input_md5().lower()!=manifest['executable_md5'].lower():
        raise ValueError('This MAP belongs to a different executable; no names were applied')
    result=dict(applied=0,already_named=0,preserved=0,unmapped=0,failed=0)
    spellings=Counter(item['name'] for item in manifest['symbols'])
    for item in manifest['symbols']:
        ea=adapter.image_base()+item['rva']
        if not adapter.is_mapped(ea):
            result['unmapped']+=1
            continue
        current=adapter.get_name(ea) or ''
        # Prefer exact C++ spelling when the receiving IDA accepts it. Otherwise
        # use the MAP token, retaining the original spelling in a repeatable comment.
        candidate=adapter.valid_name(item['name'])
        if candidate!=item['name'] or spellings[item['name']]>1:
            candidate=adapter.valid_name(item['map_name'])
        if current==candidate:
            result['already_named']+=1
            continue
        if current and not GENERIC.fullmatch(current) and not replace_existing:
            result['preserved']+=1
            continue
        if adapter.set_name(ea,candidate):
            result['applied']+=1
            if candidate!=item['name']:
                adapter.annotate(ea,'Recovered Replay name: '+item['name'])
        else:
            result['failed']+=1
    return result


class IDAAdapter:
    def __init__(self):
        import ida_bytes,ida_name,ida_nalt
        self.bytes,self.names,self.nalt=ida_bytes,ida_name,ida_nalt
    def input_md5(self):
        value=self.nalt.retrieve_input_file_md5()
        return value.hex() if isinstance(value,bytes) else str(value)
    def image_base(self):
        return self.nalt.get_imagebase()
    def is_mapped(self,ea):
        return self.bytes.is_mapped(ea)
    def get_name(self,ea):
        return self.names.get_name(ea)
    def valid_name(self,name):
        return self.names.validate_name(name,self.names.VNT_IDENT)
    def set_name(self,ea,name):
        return self.names.set_name(ea,name,self.names.SN_CHECK|self.names.SN_NOWARN)
    def annotate(self,ea,note):
        old=self.bytes.get_cmt(ea,True) or ''
        if note not in old:
            self.bytes.set_cmt(ea,(old+'\n'+note).strip(),True)


def main():
    import ida_kernwin
    path=ida_kernwin.ask_file(False,'*.map','Select the generated Replay MAP (keep its .symbols.json beside it)')
    if path:
        try:
            result=apply_package(path,IDAAdapter())
            ida_kernwin.msg('[Replay MAP] '+json.dumps(result)+'\n')
        except Exception as exc:
            ida_kernwin.warning(str(exc))


if __name__=='__main__':
    main()