"""Offline asset assembly and Spine 4.2 region export, driven by semantic v2 data."""
import argparse
import hashlib
import html
import json
from pathlib import Path

import numpy as np
from PIL import Image


def write_json(file, data):
    Path(file).write_text(json.dumps(data, ensure_ascii=False, indent=2, allow_nan=False) + '\n', encoding='utf-8')


def safe_file(root, name):
    if not isinstance(name, str) or not name or Path(name).is_absolute():
        raise ValueError(f'Expected relative asset path: {name}')
    root = Path(root).resolve()
    file = (root / name).resolve()
    if not file.is_relative_to(root) or not file.is_file():
        raise ValueError(f'Missing asset or path outside input directory: {name}')
    return file


def ordered(nodes, edges):
    """Stable topological order; do not silently break contradictory constraints."""
    pending = list(nodes)
    result = []
    while pending:
        ready = next((n for n in pending if not any(b == n and a in pending for a, b in edges)), None)
        if ready is None:
            raise ValueError('Cyclic bone or draw-order constraints')
        pending.remove(ready)
        result.append(ready)
    return result


def load_assets(input_dir, completion_dir=None, fallback_policy='sam', draw_order_file=None):
    root = Path(input_dir).resolve()
    hashes = {}

    def tracked(folder, name, label):
        file = safe_file(folder, name)
        hashes[f'{label}/{name}'] = hashlib.sha256(file.read_bytes()).hexdigest()
        return file

    def read(folder, name, label):
        return json.loads(tracked(folder, name, label).read_text(encoding='utf-8-sig'))

    source = Image.open(tracked(root, 'source.png', 'sam')).convert('RGBA')
    rgba = np.array(source)
    w, h = source.size
    plan = read(root, 'semantic-parts.json', 'sam')
    seg = read(root, 'segmentation.json', 'sam')
    parts = plan['parts']
    by_id = {p['id']: p for p in parts}
    if len(by_id) != len(parts):
        raise ValueError('Duplicate semantic ID')

    def records(items, label):
        result = {}
        for item in items:
            name = item['id']
            if name not in by_id or name in result:
                raise ValueError(f'{label}: unknown or duplicate ID: {name}')
            result[name] = item
        return result

    segmented = records(seg['parts'], 'segmentation')
    completed, skipped = {}, {}
    completion = Path(completion_dir).resolve() if completion_dir else None
    # A completion run snapshots its SAM inputs, so compare content, not directory names.
    if completion:
        for name in ['source.png', 'semantic-parts.json', 'segmentation.json']:
            tracked(completion, name, 'completion')
            if hashes[f'sam/{name}'] != hashes[f'completion/{name}']:
                raise ValueError(f'Completion input mismatch: {name}')
        result = read(completion, 'completion-result.json', 'completion')
        completed = records(result['parts'], 'completion')
        skipped = records(result.get('skipped', []), 'completion skipped')
        if set(completed) & set(skipped):
            raise ValueError('Completion ID occurs in both parts and skipped')

    images, original_images, asset_records = {}, {}, []
    masks = {}
    for name, record in segmented.items():
        if by_id[name]['visibility'] == 'hidden':
            raise ValueError(f'Hidden part unexpectedly has a SAM mask: {name}')
        mask_name = record.get('mask', f'masks/{name}.png')
        mask = np.array(Image.open(tracked(root, mask_name, 'sam')))
        if mask.shape != (h, w) or not np.isin(mask, [0, 255]).all():
            raise ValueError(f'Invalid full-canvas binary mask: {name}')
        mask = mask > 0
        if not mask.any() or np.any(mask & (rgba[:, :, 3] == 0)):
            raise ValueError(f'Empty mask or mask outside source alpha: {name}')
        if completion:
            tracked(completion, mask_name, 'completion')
            if hashes[f'sam/{mask_name}'] != hashes[f'completion/{mask_name}']:
                raise ValueError(f'Completion mask mismatch: {name}')
        masks[name] = mask

    for part in parts:
        name = part['id']
        if part['visibility'] == 'hidden':
            asset_records.append({'id': name, 'source': 'hidden', 'image': None,
                                  'reason': 'HIDDEN_NO_TEXTURE', 'semantic_verified': False})
            continue
        if name not in masks:
            raise ValueError(f'Visible part has no SAM mask: {name}')
        mask = masks[name]
        ys, xs = np.nonzero(mask)
        original_box = [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1]
        cut = rgba.copy()
        cut[~mask] = 0
        original = Image.fromarray(cut).crop(original_box)
        original_images[name] = (original, original_box)
        candidate = completed.get(name, {})
        accepted = bool(candidate.get('selected_file')) and candidate.get('qa', {}).get('accepted') is True
        warnings = list(segmented[name].get('warnings', []))
        if accepted:
            if candidate.get('status') != 'geometry_passed_needs_review':
                raise ValueError(f'Inconsistent selected completion status: {name}')
            image = Image.open(tracked(completion, candidate.get('image'), 'completion')).convert('RGBA')
            box = candidate.get('crop_xyxy')
            if (not isinstance(box, list) or len(box) != 4 or
                    any(type(x) is not int for x in box) or
                    not (0 <= box[0] < box[2] <= w and 0 <= box[1] < box[3] <= h) or
                    image.size != (box[2] - box[0], box[3] - box[1])):
                raise ValueError(f'Invalid completion crop/size: {name}')
            full = Image.new('RGBA', source.size)
            full.paste(image, (box[0], box[1]))
            if not np.array_equal(np.array(full)[mask], rgba[mask]):
                raise ValueError(f'Completion does not preserve original visible RGBA: {name}')
            origin, reason = 'completion', 'GEOMETRY_PASSED_NEEDS_REVIEW'
            if candidate.get('qa', {}).get('unfulfilled_regions'):
                warnings.append('UNFULFILLED_COMPLETION_REGIONS')
        else:
            image, box = original, original_box
            origin = 'sam'
            if not part['needs_completion']:
                reason = 'NO_COMPLETION_NEEDED'
            else:
                reason = skipped.get(name, {}).get('reason') or (
                    'COMPLETION_' + str(candidate.get('status', 'NOT_PROVIDED')).upper())
                warnings.append('USING_UNCOMPLETED_SAM')
                if fallback_policy == 'error':
                    raise ValueError(f'Completion required for {name}: {reason}')
        images[name] = (image, box)
        asset_records.append({'id': name, 'source': origin, 'reason': reason, 'image': f'images/{name}.png',
                              'crop_xyxy': box, 'original_crop_xyxy': original_box,
                              'semantic_verified': False, 'warnings': warnings,
                              'completion_status': candidate.get('status'),
                              'completion_qa': candidate.get('qa')})

    if not images:
        raise ValueError('No drawable assets')
    edges = [(p['id'], front) for p in parts for front in p['occluded_by']]
    # Keep hidden nodes while sorting so transitive constraints survive their removal.
    draw_order = [n for n in ordered(by_id, edges) if n in images]
    overridden = False
    if draw_order_file:
        file = Path(draw_order_file).resolve()
        hashes['draw_order_override'] = hashlib.sha256(file.read_bytes()).hexdigest()
        override = json.loads(file.read_text(encoding='utf-8-sig'))
        draw_order = override['draw_order']
        if not isinstance(draw_order, list) or len(draw_order) != len(images) or set(draw_order) != set(images):
            raise ValueError('draw_order must contain each drawable ID exactly once')
        overridden = True
    reach = {n: set() for n in by_id}
    for n in reversed(ordered(by_id, edges)):
        for a, b in edges:
            if a == n:
                reach[n].add(b)
                reach[n].update(reach[b])
    unresolved = []
    for i, a in enumerate(draw_order):
        for b in draw_order[i + 1:]:
            if b in reach[a] or a in reach[b]:
                continue
            ai, ab = images[a]
            bi, bb = images[b]
            x0, y0, x1, y1 = max(ab[0], bb[0]), max(ab[1], bb[1]), min(ab[2], bb[2]), min(ab[3], bb[3])
            if x0 < x1 and y0 < y1:
                am = np.array(ai.getchannel('A'))[y0-ab[1]:y1-ab[1], x0-ab[0]:x1-ab[0]] > 0
                bm = np.array(bi.getchannel('A'))[y0-bb[1]:y1-bb[1], x0-bb[0]:x1-bb[0]] > 0
                overlap = int((am & bm).sum())
                if overlap:
                    unresolved.append({'parts': [a, b], 'overlap_pixels': overlap})
    violated = [[a, b] for a in draw_order for b in reach[a]
                if b in images and draw_order.index(a) > draw_order.index(b)]
    assembly = {'schema_version': 1, 'quality_status': 'needs_review', 'semantic_verified': False,
                'width': w, 'height': h, 'coordinate_space': 'pixel_xy_top_left',
                'parts': asset_records, 'draw_order': draw_order,
                'draw_order_source': 'override' if overridden else 'occlusion_stable_topological',
                'unconstrained_overlaps': unresolved, 'overridden_occlusion_constraints': violated,
                'input_hashes': hashes}
    return source, plan, assembly, images, original_images


def build_skeleton(plan, assembly, check_angle=10):
    w, h = assembly['width'], assembly['height']
    bone_parts = {p['id']: p for p in plan['parts'] if p['rig_role'] == 'bone_part'}
    edges = [(p['parent'], name) for name, p in bone_parts.items() if p['parent']]
    bone_order = ordered(bone_parts, edges)
    world = {'root': (0, 0)}
    bones = [{'name': 'root'}]
    for name in bone_order:
        part = bone_parts[name]
        world[name] = (part['pivot_norm'][0] * w - w / 2, h - part['pivot_norm'][1] * h)
        parent = part['parent'] or 'root'
        x, y = world[name]
        px, py = world[parent]
        # Setup axes remain aligned with the canvas. Parent rotation still propagates
        # to children; attachment offsets below preserve the exact original pose.
        bones.append({'name': name, 'parent': parent, 'x': x - px, 'y': y - py})
    semantic = {p['id']: p for p in plan['parts']}
    assets = {p['id']: p for p in assembly['parts']}
    slots, attachments = [], {}
    for name in assembly['draw_order']:
        part, asset = semantic[name], assets[name]
        bone = name if part['rig_role'] == 'bone_part' else part['bind_to'] or 'root'
        bx, by = world[bone]
        x0, y0, x1, y1 = asset['crop_xyxy']
        slots.append({'name': name, 'bone': bone, 'attachment': name})
        attachments[name] = {name: {'type': 'region', 'path': name,
                                   'x': (x0 + x1) / 2 - w / 2 - bx,
                                   'y': h - (y0 + y1) / 2 - by,
                                   'width': x1 - x0, 'height': y1 - y0}}
    animations = {'setup': {}}
    for name in bone_order:
        animations[f'check__{name}'] = {'bones': {name: {'rotate': [
            {'time': 0, 'value': 0}, {'time': .5, 'value': check_angle},
            {'time': 1, 'value': 0}, {'time': 1.5, 'value': -check_angle}, {'time': 2, 'value': 0}]}}}
    return {'skeleton': {'spine': '4.2.0', 'x': -w / 2, 'y': 0, 'width': w, 'height': h,
                         'images': './images/'},
            'bones': bones, 'slots': slots, 'skins': [{'name': 'default', 'attachments': attachments}],
            'animations': animations}


def compose(size, images, order):
    canvas = Image.new('RGBA', size)
    coverage = np.zeros((size[1], size[0]), np.uint16)
    for name in order:
        image, (x0, y0, x1, y1) = images[name]
        canvas.alpha_composite(image, (x0, y0))
        coverage[y0:y1, x0:x1] += np.array(image.getchannel('A')) > 0
    return canvas, coverage


def export_assets(data, output, check_angle=10):
    source, plan, assembly, images, original_images = data
    out = Path(output)
    (out / 'images').mkdir(parents=True, exist_ok=True)
    source.save(out / 'source.png')
    write_json(out / 'semantic-parts.json', plan)
    for name, (im, _) in images.items():
        im.save(out / 'images' / f'{name}.png')
    assembled, coverage = compose(source.size, images, assembly['draw_order'])
    before, _ = compose(source.size, original_images, assembly['draw_order'])
    assembled.save(out / 'assembled.png')
    before.save(out / 'assembled-sam.png')
    foreground = np.array(source.getchannel('A')) > 0
    missing = foreground & (coverage == 0)
    diagnostic = np.zeros((source.height, source.width, 4), np.uint8)
    diagnostic[coverage > 1] = [255, 180, 0, 220]
    diagnostic[missing] = [255, 0, 80, 255]
    Image.fromarray(diagnostic).save(out / 'coverage-diagnostic.png')
    comparison = Image.new('RGBA', (source.width * 3, source.height), 'white')
    for i, im in enumerate([source, before, assembled]):
        comparison.alpha_composite(im, (i * source.width, 0))
    comparison.convert('RGB').save(out / 'comparison.png')
    assembly['diagnostics'] = {'uncovered_source_pixels': int(missing.sum()),
                               'overlap_pixels': int((coverage > 1).sum()),
                               'rotation_quality_verified': False}
    write_json(out / 'assembly.json', assembly)
    skeleton = build_skeleton(plan, assembly, check_angle)
    write_json(out / 'skeleton.json', skeleton)
    rows = []
    for part in assembly['parts']:
        picture = f'<img src="{part["image"]}">' if part['image'] else ''
        rows.append('<tr><td>' + html.escape(part['id']) + '</td><td>' + html.escape(part['source']) +
                    '</td><td>' + html.escape(part['reason']) + '</td><td>' + picture + '</td></tr>')
    notes = {key: assembly[key] for key in ['diagnostics', 'draw_order_source', 'unconstrained_overlaps',
                                           'overridden_occlusion_constraints']}
    (out / 'report.html').write_text('''<!doctype html><html lang="zh-CN"><meta charset="utf-8">
<title>第四阶段：Spine 组装与导出</title>
<style>body{font:16px system-ui;margin:24px}img{max-width:100%;background:#ddd}td img{max-height:150px}td,th{border:1px solid #ccc;padding:8px}pre{white-space:pre-wrap}</style>
<h1>第四阶段：Spine 组装与导出</h1>
<p>状态：needs_review。几何通过与成功导出不代表语义、接缝或动作质量合格。</p>
<p><a href="preview.html">Spine 预览</a>（播放器从 CDN 加载，需要网络）；选择 setup 查看静态姿态，选择 check__骨骼ID 检查单关节旋转及子骨骼跟随。检查动画不代表成品动作。</p>
<p><a href="skeleton.json">骨架 JSON</a> · <a href="assembly.json">来源与诊断</a></p>
<h2>原图 / SAM 组装 / 最终组装</h2><img src="comparison.png">
<h2>覆盖诊断：红色为原图漏分，黄色为部件重叠</h2><img src="coverage-diagnostic.png">
<p>重叠可能是正常遮挡；覆盖率不能判断语义。没有遮挡约束的重叠使用语义计划中的稳定顺序，需复核。</p>
<pre>''' + html.escape(json.dumps(notes, ensure_ascii=False, indent=2)) + '''</pre>
<h2>资产来源</h2><table><tr><th>ID</th><th>来源</th><th>原因</th><th>部件</th></tr>''' +
                                    ''.join(rows) + '</table></html>', encoding='utf-8')
    return {'parts': len(images), 'bones': len(skeleton['bones']),
            'completion_parts': sum(p['source'] == 'completion' for p in assembly['parts']),
            'fallback_parts': sum('USING_UNCOMPLETED_SAM' in p.get('warnings', []) for p in assembly['parts']),
            'quality_status': 'needs_review', 'input_hashes': assembly['input_hashes']}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input-dir', required=True)
    parser.add_argument('--completion-dir')
    parser.add_argument('--output')
    parser.add_argument('--draw-order-file')
    parser.add_argument('--fallback-policy', choices=['sam', 'error'], default='sam')
    parser.add_argument('--check-angle', type=float, default=10)
    args = parser.parse_args()
    if not 0 < args.check_angle <= 45:
        parser.error('check-angle must be in (0, 45]')
    data = load_assets(args.input_dir, args.completion_dir, args.fallback_policy, args.draw_order_file)
    if args.output:
        result = export_assets(data, args.output, args.check_angle)
    else:
        skeleton = build_skeleton(data[1], data[2], args.check_angle)
        result = {'input_hashes': data[2]['input_hashes'], 'parts': len(data[3]), 'bones': len(skeleton['bones'])}
    print(json.dumps(result, ensure_ascii=False))


if __name__ == '__main__':
    main()