"""Adaptive completion planning, reference canvases and reversible native-pixel merging."""
import hashlib
import json
import math
from pathlib import Path

import cv2
import numpy as np
from PIL import Image, ImageDraw

from normalize_reference import prepare_rgba


def read(file):
    return json.loads(Path(file).read_text(encoding='utf-8-sig'))


def write(file, value):
    Path(file).write_text(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False), encoding='utf-8')


def binary(file, shape):
    a = np.array(Image.open(file))
    if a.shape != shape or not np.isin(a, [0, 255]).all():
        raise ValueError(f'Expected full-canvas binary PNG: {file}')
    return a > 0


def review_inputs(file, ids, shape):
    """Review masks are explicitly authored data; never guess ownership from scores."""
    if not file:
        return {}, {}
    file = Path(file).resolve()
    data = read(file)
    if data.get('schema_version') != 1 or not isinstance(data.get('parts'), dict):
        raise ValueError('Completion review requires schema_version=1 and parts object')
    hashes = {'review.json': hashlib.sha256(file.read_bytes()).hexdigest()}
    result = {}
    for name, entry in data['parts'].items():
        if name not in ids or not isinstance(entry, dict) or entry.get('confirmed') is not True:
            raise ValueError(f'Review entry requires known ID and confirmed=true: {name}')
        if set(entry) - {'confirmed', 'extension_radius', 'repair_mask', 'replace_mask', 'remove_mask'}:
            raise ValueError(f'Unknown review fields: {name}')
        item = {'confirmed': True}
        if 'extension_radius' in entry:
            radius = entry['extension_radius']
            if type(radius) is not int or not 1 <= radius <= 256:
                raise ValueError(f'Invalid reviewed radius: {name}')
            item['extension_radius'] = radius
        for key in ['repair_mask', 'replace_mask', 'remove_mask']:
            if key not in entry:
                continue
            ref = entry[key]
            if not isinstance(ref, str) or Path(ref).is_absolute():
                raise ValueError('Review mask must be relative to review file')
            mask_file = (file.parent / ref).resolve()
            if not mask_file.is_relative_to(file.parent):
                raise ValueError('Review mask path escapes review directory')
            item[key] = binary(mask_file, shape)
            hashes[f'{name}/{key}'] = hashlib.sha256(mask_file.read_bytes()).hexdigest()
        result[name] = item
    return result, hashes


def bounding(mask):
    ys, xs = np.nonzero(mask)
    return [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1] if len(xs) else None


def enclosed(mask):
    padded = np.pad(mask.astype(np.uint8), 1)
    fill = padded.copy()
    cv2.floodFill(fill, None, (0, 0), 1)
    return ~fill[1:-1, 1:-1].astype(bool)


def white_image(array):
    im = Image.fromarray(array).convert('RGBA')
    return Image.alpha_composite(Image.new('RGBA', im.size, 'white'), im).convert('RGB')


def prepare_adaptive(directory, selected_ids, size=1024, radius=None, reference_mode='board', review_file=None):
    # Import shared native coordinate helpers lazily to avoid a module cycle.
    from completion_geometry import localize, allowed_regions
    out = Path(directory)
    plan, seg = read(out / 'semantic-parts.json'), read(out / 'segmentation.json')
    rgba = np.array(Image.open(out / 'source.png').convert('RGBA'))
    h, w = rgba.shape[:2]
    ids = {p['id']: p for p in plan['parts']}
    if set(selected_ids) - set(ids):
        raise ValueError('Unknown requested part ID')
    records = {p['id']: p for p in seg['parts']}
    if len(records) != len(seg['parts']) or set(records) - set(ids):
        raise ValueError('Invalid segmentation IDs')
    masks = {name: binary(out / 'masks' / f'{name}.png', (h, w)) for name in records}
    for name, mask in masks.items():
        if not mask.any() or np.any(mask & (rgba[:, :, 3] == 0)):
            raise ValueError(f'Empty mask or mask outside original alpha: {name}')
    reviews, review_hashes = review_inputs(review_file, ids, (h, w))
    coverage = np.zeros((h, w), np.uint16)
    for mask in masks.values():
        coverage += mask
    tasks, skipped, assessments = [], [], []
    (out / 'preservation').mkdir(exist_ok=True)
    for part in plan['parts']:
        name = part['id']
        if selected_ids and name not in selected_ids:
            continue
        if name not in masks or part['visibility'] == 'hidden':
            skipped.append({'id': name, 'reason': 'NO_VISIBLE_SAM_MASK'})
            continue
        visible = masks[name]
        box = bounding(visible)
        review = reviews.get(name, {})
        resolved_radius = review.get('extension_radius', radius if radius is not None else
                                     int(np.clip(math.ceil(min(box[2]-box[0], box[3]-box[1]) * .15), 4, 96)))
        inferred = dict(part)
        inferred['completion'] = {
            'occlusion_fill': bool(part['occluded_by']),
            'joint_extension_parent': part['rig_role'] == 'bone_part' and part['parent'] is not None,
            'joint_extension_child': part['rig_role'] == 'bone_part' and any(
                p['parent'] == name and p['rig_role'] == 'bone_part' for p in plan['parts'])}
        _, regions, warnings = allowed_regions(inferred, plan, masks, visible, resolved_radius)
        # Only enclosed source foreground without another mask owner is an automatic
        # repair candidate. Open gaps and ambiguous ownership stay diagnostic.
        gap = enclosed(visible) & (rgba[:, :, 3] > 0) & (coverage == 0)
        cleanup = np.zeros_like(visible)
        for front in part['occluded_by']:
            p = ids[front]
            if front in masks and p['rig_role'] != 'bone_part' and p['bind_to'] == name:
                cleanup |= visible & masks[front]
        replace = review.get('replace_mask', np.zeros_like(visible))
        removed = review.get('remove_mask', np.zeros_like(visible))
        repair = review.get('repair_mask', np.zeros_like(visible))
        if np.any((replace | removed) & ~visible) or np.any(repair & visible) or np.any(replace & removed):
            raise ValueError(f'{name}: replacement/removal must be disjoint subsets of SAM; repair must be outside SAM')
        regions['segmentation_gap'] = gap | repair
        regions['attachment_cleanup'] = cleanup & ~removed
        regions['reviewed_replacement'] = replace
        allowed = np.zeros_like(visible)
        for region in regions.values():
            allowed |= region
        allowed &= ~removed
        protected = visible & ~allowed & ~removed
        unexpected = []
        duplicates = []
        for other, mask in masks.items():
            if other == name:
                continue
            overlap = int((mask & visible).sum())
            related = other in part['occluded_by'] or name in ids[other]['occluded_by']
            if overlap and not related:
                unexpected.append({'id': other, 'pixels': overlap})
                if overlap / max(1, int((mask | visible).sum())) >= .8:
                    duplicates.append(other)
        assessment = {'id': name, 'original_needs_completion': part['needs_completion'],
                      'inferred_completion': inferred['completion'], 'radius': resolved_radius,
                      'region_pixels': {k: int(v.sum()) for k, v in regions.items()},
                      'protected_pixels': int(protected.sum()), 'removed_pixels': int(removed.sum()),
                      'unexpected_overlaps': unexpected, 'review_confirmed': bool(review)}
        assessments.append(assessment)
        if duplicates and not review:
            skipped.append({'id': name, 'reason': 'AMBIGUOUS_MASK_REQUIRES_REVIEW', 'overlaps': duplicates})
            continue
        if not protected.any() or not (protected & (rgba[:, :, 3] >= 240)).any():
            skipped.append({'id': name, 'reason': 'NO_PRESERVED_ALIGNMENT_ANCHOR'})
            continue
        if not allowed.any() and not removed.any():
            skipped.append({'id': name, 'reason': 'NO_OBSERVED_EDIT_REGION'})
            continue
        region_box = bounding(visible | allowed)
        bw, bh = region_box[2] - region_box[0], region_box[3] - region_box[1]
        native_size = max(8, math.ceil(max(bw, bh) * 1.15))
        origin = [region_box[0] - (native_size-bw)//2, region_box[1] - (native_size-bh)//2]
        work_size = int(size * .68) if reference_mode == 'board' else size
        offset = [int(size * .02), (size-work_size)//2] if reference_mode == 'board' else [0, 0]
        scale = work_size / native_size
        folder = out / 'tasks' / name
        folder.mkdir(parents=True, exist_ok=True)
        native_source = localize(rgba, origin, native_size)
        native_visible = localize(visible, origin, native_size)
        native_keep = localize(protected, origin, native_size)
        native_edit = localize(allowed, origin, native_size)
        native_remove = localize(removed, origin, native_size)
        cut = native_source.copy()
        cut[~native_keep] = 0
        Image.fromarray(cut).save(folder / 'native-visible.png')
        Image.fromarray(native_source).save(folder / 'native-source.png')
        for label, mask in [('keep', native_keep), ('allowed', native_edit), ('remove', native_remove), ('sam', native_visible)]:
            Image.fromarray(mask.astype(np.uint8)*255).save(folder / f'native-{label}.png')

        def project(array, is_mask=False):
            small = np.array(Image.fromarray(array.astype(np.uint8)*255 if is_mask else array).resize(
                (work_size, work_size), Image.Resampling.NEAREST if is_mask else Image.Resampling.LANCZOS))
            canvas = np.zeros((size, size) + small.shape[2:], np.uint8)
            canvas[offset[1]:offset[1]+work_size, offset[0]:offset[0]+work_size] = small
            return canvas

        model_visible = project(cut)
        Image.fromarray(model_visible).save(folder / 'visible.png')
        model_allowed = project(native_edit, True) > 0
        model_remove = project(native_remove, True) > 0
        Image.fromarray(model_allowed.astype(np.uint8)*255).save(folder / 'allowed-mask.png')
        mask_rgba = np.full((size, size, 4), 255, np.uint8)
        mask_rgba[model_allowed | model_remove, 3] = 0
        Image.fromarray(mask_rgba).save(folder / 'edit-mask.png')
        context = white_image(project(native_source))
        context.save(folder / 'context.png')
        full_reference = white_image(rgba)
        full_reference.thumbnail((size, size), Image.Resampling.LANCZOS)
        full_reference.save(folder / 'reference.png')
        model_input = white_image(model_visible)
        panel_boxes = []
        if reference_mode == 'board':
            draw = ImageDraw.Draw(model_input)
            x = int(size * .74)
            draw.line([(x-4, 0), (x-4, size)], fill='#777777', width=1)
            for label, image, top in [('LOCAL REFERENCE', white_image(native_source), 0),
                                      ('FULL REFERENCE', full_reference.copy(), size//2)]:
                draw.text((x, top+2), label, fill='black')
                image.thumbnail((size-x-4, size//2-24), Image.Resampling.LANCZOS)
                model_input.paste(image, (x, top+22))
                panel_boxes.append([x, top+22, x+image.width, top+22+image.height])
        model_input.save(folder / 'input.png')
        bounds, native_counts = {}, {}
        for kind, region in regions.items():
            native = localize(region & allowed, origin, native_size)
            model = project(native, True)
            Image.fromarray(native.astype(np.uint8)*255).save(folder / f'native-{kind}.png')
            Image.fromarray(model).save(folder / f'{kind}.png')
            bounds[kind] = bounding(model > 0)
            native_counts[kind] = int(native.sum())
        preservation = {}
        for key, mask in [('preserve_mask', protected), ('editable_mask', allowed), ('removed_mask', removed)]:
            relative = f'preservation/{name}-{key}.png'
            Image.fromarray(mask.astype(np.uint8)*255).save(out / relative)
            preservation[key] = relative
        if unexpected:
            warnings.append('UNEXPECTED_OVERLAP_REVIEW_REQUIRED')
        if cleanup.any():
            warnings.append('ATTACHMENT_BASE_CLEANUP_NEEDS_REVIEW')
        task = {'id': name, 'semantic': part['semantic'], 'profile': 'adaptive',
                'completion': inferred['completion'], 'original_completion': part['completion'],
                'origin_xy': origin, 'native_size': native_size, 'canvas_size': [size, size], 'source_size': [w, h],
                'extension_radius': resolved_radius, 'allowed_pixels': int(native_edit.sum()),
                'removed_pixels': int(native_remove.sum()), 'regions': native_counts,
                'region_bounds_xyxy': bounds, 'occluded_by': part['occluded_by'],
                'warnings': records[name].get('warnings', []) + warnings, 'preservation': preservation,
                'reference_mode': reference_mode, 'reference_panel_boxes': panel_boxes,
                'transform': {'scale': scale, 'working_xyxy': [*offset, offset[0]+work_size, offset[1]+work_size],
                              'source_to_canvas': [scale, 0, offset[0]-scale*origin[0]+(scale-1)/2,
                                                   0, scale, offset[1]-scale*origin[1]+(scale-1)/2],
                              'canvas_to_source': [1/scale, 0, origin[0]-offset[0]/scale+(1/scale-1)/2,
                                                   0, 1/scale, origin[1]-offset[1]/scale+(1/scale-1)/2]},
                'semantic_verified': False, 'status': 'prepared'}
        write(folder / 'task.json', task)
        tasks.append(task)
    result = {'schema_version': 2, 'profile': 'adaptive', 'tasks': tasks, 'skipped': skipped,
              'assessments': assessments, 'review_hashes': review_hashes, 'reference_mode': reference_mode,
              'canvas_size': size, 'quality_status': 'needs_review'}
    write(out / 'completion-plan.json', result)
    return result


def merge_adaptive(folder, candidate_file, output_file, size_policy):
    folder = Path(folder)
    task = read(folder / 'task.json')
    candidate = Image.open(candidate_file).convert('RGBA')
    qa = {'accepted': False, 'semantic_verified': False, 'reasons': [], 'profile': 'adaptive',
          'expected_size': task['canvas_size'], 'returned_size': list(candidate.size),
          'size_policy': size_policy, 'canvas_resampled': False, 'transform': task['transform']}
    tw, th = task['canvas_size']
    if candidate.size != (tw, th):
        if size_policy == 'strict' or candidate.width * th != candidate.height * tw:
            qa['reasons'].append('CANVAS_SIZE_MISMATCH' if size_policy == 'strict' else 'CANVAS_ASPECT_RATIO_MISMATCH')
            return qa
        qa.update(canvas_resampled=True, canvas_scale=tw/candidate.width, resampling='LANCZOS')
        candidate = candidate.resize((tw, th), Image.Resampling.LANCZOS)
        candidate.save(Path(output_file).with_name(Path(output_file).stem + '-normalized.png'))
    # Record raw behavior before masking/restoration; an accepted HTTP response is
    # not evidence that the relay/model respected the supplied edit mask.
    expected_input = np.array(Image.open(folder / 'input.png').convert('RGB')).astype(float)
    raw = np.array(white_image(np.array(candidate))).astype(float)
    editable = np.array(Image.open(folder / 'edit-mask.png'))[:, :, 3] == 0
    outside_delta = np.abs(raw - expected_input).mean(axis=2)
    qa['raw_outside_edit_mae'] = float(outside_delta[~editable].mean()) if (~editable).any() else 0
    qa['raw_outside_edit_changed_pixels'] = int(((outside_delta > 30) & ~editable).sum())
    rect = task['transform']['working_xyxy']
    try:
        cropped, preprocessing = prepare_rgba(candidate.crop(rect))
    except ValueError as error:
        qa.update(reasons=['CANNOT_EXTRACT_FOREGROUND'], detail=str(error))
        return qa
    side = task['native_size']
    generated = np.array(cropped.resize((side, side), Image.Resampling.LANCZOS))
    original = np.array(Image.open(folder / 'native-visible.png').convert('RGBA'))
    source = np.array(Image.open(folder / 'native-source.png').convert('RGBA'))
    keep = binary(folder / 'native-keep.png', (side, side))
    allowed = binary(folder / 'native-allowed.png', (side, side))
    removed = binary(folder / 'native-remove.png', (side, side))
    # Judge alignment before the inverse resampling. Alpha at native boundary
    # pixels changes during interpolation even for an exactly aligned candidate.
    model_original = np.array(Image.open(folder / 'visible.png').convert('RGBA').crop(rect))
    model_generated = np.array(cropped)
    opaque = model_original[:, :, 3] >= 240
    if not opaque.any():
        qa['reasons'].append('NO_OPAQUE_ALIGNMENT_ANCHOR')
        return qa
    rgb = np.array(white_image(model_generated)).astype(float)
    mae = float(np.abs(rgb[opaque] - model_original[:, :, :3][opaque]).mean())
    retained = float((model_generated[:, :, 3][opaque] >= 240).mean())
    proposed = (generated[:, :, 3] > 0) & ~keep
    additions = proposed & allowed
    _, labels = cv2.connectedComponents((keep | additions).astype(np.uint8), connectivity=8)
    touching = np.unique(labels[keep])
    connected = additions & np.isin(labels, touching[touching != 0])
    counts = {kind: int((connected & binary(folder / f'native-{kind}.png', (side, side))).sum())
              for kind in task['regions']}
    qa.update(alignment_mae=mae, visible_retention=retained, alignment_space='request_work_panel', new_pixels=int(connected.sum()),
              removed_pixels=int(removed.sum()), region_additions=counts,
              unfulfilled_regions=[k for k, n in task['regions'].items() if n and not counts[k]],
              discarded_outside_allowed=int((proposed & ~allowed).sum()), preprocessing=preprocessing)
    if mae > 30 or retained < .95:
        qa['reasons'].append('VISIBLE_ALIGNMENT_MISMATCH')
    if not connected.any() and not removed.any():
        qa['reasons'].append('NO_CONNECTED_NEW_PIXELS')
    for kind in ['attachment_cleanup', 'reviewed_replacement']:
        if task['regions'][kind] and counts[kind] / task['regions'][kind] < .95:
            qa['reasons'].append(f'INCOMPLETE_{kind.upper()}')
    if qa['reasons']:
        return qa
    merged = np.zeros_like(original)
    merged[connected] = generated[connected]
    # Visible segmentation gaps already have ground-truth texture in source.
    # Restore it where the proposed fill connects, rather than trusting a repaint.
    repair = binary(folder / 'native-segmentation_gap.png', (side, side)) & connected & (source[:, :, 3] > 0)
    merged[repair] = source[repair]
    merged[keep] = original[keep]
    merged[removed] = 0
    qa.update(accepted=True, quality_status='needs_review', visible_pixels_exact=True,
              preservation_scope='preserve_mask', repaired_source_pixels=int(repair.sum()))
    Image.fromarray(merged).save(output_file)
    return qa