"""Measured mask QA and previews made only from the original character pixels."""
import colorsys
import html
import json
import math
import re
from pathlib import Path

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


def assess_candidate(mask, alpha, bbox, points):
    mask = np.asarray(mask, dtype=bool)
    if mask.shape != alpha.shape:
        raise ValueError('Mask dimensions do not match the source')
    clipped = mask & (alpha > 0)
    _, _, stats, _ = cv2.connectedComponentsWithStats(clipped.astype(np.uint8), 8)
    area, foreground_area = int(mask.sum()), int(clipped.sum())
    x0, y0, x1, y1 = bbox
    hit = sum(bool(clipped[y, x]) for x, y in points)
    return {
        'raw_area_px': area, 'area_px': foreground_area,
        'outside_alpha_ratio': (area - foreground_area) / max(1, area),
        'inside_roi_fraction': int(clipped[y0:y1, x0:x1].sum()) / max(1, foreground_area),
        'seed_hit_fraction': hit / max(1, len(points)),
        'components_at_least_4px': int((stats[1:, cv2.CC_STAT_AREA] >= 4).sum()),
    }


def render_outputs(image, parts, out_dir):
    out = Path(out_dir)
    (out / 'parts').mkdir(parents=True, exist_ok=True)
    (out / 'masks').mkdir(parents=True, exist_ok=True)
    rgba = np.array(image.convert('RGBA'))
    h, w = rgba.shape[:2]
    coverage = np.zeros((h, w), np.uint16)
    white = Image.new('RGBA', (w, h), 'white')
    overlay = Image.alpha_composite(white, image.convert('RGBA'))
    records, crops, clipped_masks = [], [], []
    seen = set()
    for index, part in enumerate(parts):
        name = part['id']
        if not re.fullmatch('[a-z][a-z0-9_-]{0,63}', name) or name in seen:
            raise ValueError(f'Invalid or duplicate part ID: {name}')
        seen.add(name)
        mask = np.asarray(part['mask'], dtype=bool)
        if mask.shape != (h, w):
            raise ValueError(f'{name}: mask dimensions mismatch')
        mask = mask & (rgba[:, :, 3] > 0)
        clipped_masks.append(mask)
        ys, xs = np.nonzero(mask)
        if not len(xs):
            raise ValueError(f'{name}: empty foreground mask')
        coverage += mask
        bbox = [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1]
        color = [round(v * 255) for v in colorsys.hsv_to_rgb((index * .61803398875) % 1, .75, .95)]
        cut = rgba.copy()
        cut[~mask] = 0
        crop = Image.fromarray(cut).crop(bbox)
        crop.save(out / 'parts' / f'{name}.png')
        Image.fromarray(mask.astype(np.uint8) * 255).save(out / 'masks' / f'{name}.png')
        tint = np.zeros_like(rgba)
        tint[mask, :3] = color
        tint[mask, 3] = 100
        overlay = Image.alpha_composite(overlay, Image.fromarray(tint))
        records.append({k: v for k, v in part.items() if k != 'mask'} | {
            'color_rgb': color, 'crop_xyxy': bbox, 'area_px': int(mask.sum()),
            'image': f'parts/{name}.png', 'mask': f'masks/{name}.png',
        })
        crops.append(crop)
    near_duplicates = []
    for i, a in enumerate(clipped_masks):
        for j in range(i + 1, len(clipped_masks)):
            b = clipped_masks[j]
            iou = int((a & b).sum()) / max(1, int((a | b).sum()))
            if iou >= .8:
                near_duplicates.append({'part_a': records[i]['id'], 'part_b': records[j]['id'], 'iou': iou})
                for index in [i, j]:
                    warnings = records[index].setdefault('warnings', []).copy()
                    if 'NEAR_DUPLICATE_MASK' not in warnings:
                        warnings.append('NEAR_DUPLICATE_MASK')
                    records[index]['warnings'] = warnings
    # Keep native scale, separated by padding and a label band. No rotation or
    # resampling of source pixels is used in the exploded sheet.
    pad, label = 24, 26
    cell_widths = [max(c.width, len(r['id']) * 8 + 8) for c, r in zip(crops, records)]
    sheet_width = max(640, max(cell_widths, default=1) + pad * 2,
                      int(math.sqrt(sum((cw + pad) * (c.height + label + pad) for cw, c in zip(cell_widths, crops)) * 1.5)))
    x, y, row_height = pad, pad, 0
    for record, crop, cw in zip(records, crops, cell_widths):
        if x + cw + pad > sheet_width and x != pad:
            x, y, row_height = pad, y + row_height + pad, 0
        record['sheet_xyxy'] = [x, y + label, x + crop.width, y + label + crop.height]
        x += cw + pad
        row_height = max(row_height, label + crop.height)
    sheet = Image.new('RGB', (sheet_width, max(100, y + row_height + pad)), 'white')
    draw = ImageDraw.Draw(sheet)
    font = ImageFont.load_default(size=14)
    for record, crop in zip(records, crops):
        x0, y0, _, _ = record['sheet_xyxy']
        draw.rectangle((x0, y0 - label + 4, x0 + 10, y0 - label + 14), fill=tuple(record['color_rgb']))
        draw.text((x0 + 16, y0 - label + 1), record['id'], fill='black', font=font)
        sheet.paste(crop, (x0, y0), crop)
    overlay.convert('RGB').save(out / 'annotated.png')
    sheet.save(out / 'parts-sheet.png')
    Image.fromarray(((rgba[:, :, 3] > 0) & (coverage == 0)).astype(np.uint8) * 255).save(out / 'unassigned-mask.png')
    Image.fromarray((coverage > 1).astype(np.uint8) * 255).save(out / 'overlap-mask.png')
    total = int((rgba[:, :, 3] > 0).sum())
    return {'parts': records, 'near_duplicate_masks': near_duplicates, 'coverage': {
        'foreground_pixels': total, 'covered_pixels': int((coverage > 0).sum()),
        'coverage_ratio': int((coverage > 0).sum()) / max(1, total),
        'overlap_pixels': int((coverage > 1).sum()),
        'unassigned_pixels': int(((rgba[:, :, 3] > 0) & (coverage == 0)).sum()),
    }, 'outputs': {'annotated': 'annotated.png', 'parts_sheet': 'parts-sheet.png'}}


def write_report(out_dir, result):
    out = Path(out_dir)
    rows = ''.join(f'<tr><td><span style="background:rgb({",".join(map(str, p["color_rgb"]))})">&nbsp;&nbsp;</span> <a href="{p["image"]}">{html.escape(p["id"])}</a></td><td>{html.escape(p["semantic"])}</td><td>{p["model_score"]:.3f}</td><td>{html.escape(", ".join(p.get("warnings", [])))}</td></tr>' for p in result['parts'])
    body = f'''<!doctype html><html lang="zh-CN"><meta charset="utf-8"><title>SAM 部件分割</title>
<style>body{{font:16px system-ui;max-width:1400px;margin:24px auto;padding:16px}}img{{max-width:100%;height:auto}}td,th{{border:1px solid #ccc;padding:8px}}table{{border-collapse:collapse}}pre{{white-space:pre-wrap}}</style>
<h1>SAM 本地分割结果</h1><p>颜色表示候选部件 mask。白底拆解图保留原始可见像素，尚未补全遮挡部分。重叠 mask 保留并计入诊断，不代表已经获得互斥、语义正确的最终拆件。</p>
<h2>原图半透明标注</h2><img src="annotated.png"><h2>白底拆解图</h2><img src="parts-sheet.png">
<h2>部件</h2><table><tr><th>ID / 颜色</th><th>语义</th><th>模型分数</th><th>需复核</th></tr>{rows}</table>
<h2>未生成的部件</h2><pre>{html.escape(json.dumps(result['skipped'], ensure_ascii=False, indent=2))}</pre>
<h2>覆盖与重叠</h2><p>覆盖率只表示有 mask 的像素比例，不是分割准确率。模型分数也不代表部件语义已验证。</p><p><a href="unassigned-mask.png">漏分区域（白色）</a> · <a href="overlap-mask.png">重叠区域（白色）</a></p><pre>{html.escape(json.dumps(result['coverage'], indent=2))}</pre>
<h2>疑似重复分割</h2><p>以下不同部件的 mask IoU ≥ 0.8，需要检查是否分到了同一物体。</p><pre>{html.escape(json.dumps(result['near_duplicate_masks'], indent=2))}</pre></html>'''
    (out / 'report.html').write_text(body, encoding='utf-8')