"""Experimental, non-destructive spread splitting with an HTML review gallery.

python -m copienator.auto_page_split INPUT OUTPUT [--workers 4]
No booklet reordering, secondary cuts, or GUI changes. Unknown formats and
uncertain content are retained. OUTPUT must not already exist.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ProcessPoolExecutor
import hashlib
import html
import json
import multiprocessing
from pathlib import Path
from urllib.parse import quote

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

from copienator.ink_detection import detect_bounds, detect_page_content


def classify_format(width: float, height: float) -> str:
    """Dimensions are displayed PDF points; tolerate 5% scan size error."""
    w, h = width * 25.4 / 72, height * 25.4 / 72
    if abs(min(w, h) / 210 - 1) < .05 and abs(max(w, h) / 297 - 1) < .05:
        return 'a4'
    if abs(w / 420 - 1) < .05 and abs(h / 297 - 1) < .05:
        return 'a3'
    return 'unknown'


def analyse_ink(rgb: np.ndarray, dpi: float, *, spread: bool = False) -> tuple[np.ndarray, dict]:
    """Use normal rognage on logical pages; colour cleanup only for overflow.

    Blankness uses the shared crop detector's content-analysis entry point.
    No independent weak-pixel veto or secondary artifact detector is applied.
    """
    if not spread:
        return detect_page_content(rgb,dpi)
    masks = {}
    result = detect_bounds(rgb, dpi, content_masks=masks,
                           colour_cleanup=True, edge_mark_cleanup=False)
    ink = masks['foreground']
    # For ownership, follow connected faint pen strokes beyond rognage's local
    # 2 mm recovery radius. Do not restore page-wide connected ruling networks.
    candidate = ink | masks['weak_candidates']
    n,labels,stats,_ = cv2.connectedComponentsWithStats(candidate.astype(np.uint8),8)
    seeded = np.zeros(n,bool)
    seeded[np.unique(labels[ink])] = True
    seeded[0] = False
    giant = (stats[:,2] > .65*ink.shape[1]) & (stats[:,3] > .4*ink.shape[0])
    seeded[giant] = False
    ink = ink | seeded[labels]
    return ink, result


def _analyse_page_full(rgb: np.ndarray, dpi: float, kind: str) -> dict:
    """Keep spread geometry analysis separate from logical-page blankness."""
    h, w = rgb.shape[:2]
    boundaries = [0, w//2, w] if kind == 'a3' else [0, w]
    content = np.zeros((h,w), bool)
    evidence = []
    for x0,x1 in zip(boundaries,boundaries[1:]):
        foreground, result = analyse_ink(rgb[:,x0:x1],dpi)
        content[:,x0:x1] = foreground
        evidence.append(result)
    ink = analyse_ink(rgb,dpi,spread=True)[0] if kind == 'a3' else content
    plan = plan_masks(ink,dpi,kind,content_ink=content,rgb=rgb)
    for output,result in zip(plan['outputs'],evidence):
        output['crop_analysis'] = result
    return plan


def analyse_page(rgb: np.ndarray, dpi: float, kind: str) -> dict:
    """Analyse a raster with the standard coarse-to-fine splitter."""
    from copienator.fast_page_split import analyse_page_fast
    return analyse_page_fast(rgb, dpi, kind)


def _fold_regions(ink: np.ndarray, credible: np.ndarray, dpi: float,
                  rgb: np.ndarray | None = None) -> tuple[list[dict], np.ndarray]:
    """Grow local word groups towards page anchors, stopping at supported text.

    Broader gaps are allowed for ownership evidence without merging both pages
    into one giant line or changing output bounds to include the anchor word.
    """
    import heapq
    px = dpi/25.4
    h,w = ink.shape
    middle = w//2
    ignored = []
    if rgb is not None:
        # Remove weak thin ruling *before* word grouping: otherwise a nearby
        # dark word can lend its colour/anchor evidence to a pale grid fragment.
        count,raw_labels,raw_stats,_ = cv2.connectedComponentsWithStats(ink.astype(np.uint8),8)
        ink = ink.copy()
        for i,(x,y,bw,bh,area) in enumerate(raw_stats[1:],1):
            if not (x < middle+2*px and x+bw > middle-2*px and bh<px and bw>4*bh):
                continue
            region = raw_labels[y:y+bh,x:x+bw]==i
            values = rgb[y:y+bh,x:x+bw][region].astype(np.int16)
            if np.median(np.ptp(values,axis=1))<=80 and np.median(values.min(axis=1))>=110:
                ink[y:y+bh,x:x+bw][region] = False
                ignored.append(dict(owner='artifact',box=[int(x),int(y),int(x+bw),int(y+bh)],
                    component_id=0,anchor_boxes=[None,None],anchor_distance_mm=[None,None],reason='weak-thin-ruling'))
    kernel = np.ones((max(1,round(.5*px)),max(1,round(3*px))),np.uint8)
    joined = cv2.dilate(ink.astype(np.uint8),kernel)
    n,labels,stats,_ = cv2.connectedComponentsWithStats(joined,8)
    boxes = np.zeros((n,4),np.int32)
    areas = np.bincount(labels[ink],minlength=n)
    pen = np.zeros(n,bool)
    for i in range(1,n):
        x,y,bw,bh,_ = stats[i]
        yy,xx = np.nonzero((labels[y:y+bh,x:x+bw]==i) & ink[y:y+bh,x:x+bw])
        boxes[i] = [x+xx.min(),y+yy.min(),x+xx.max()+1,y+yy.max()+1]
        if rgb is not None and len(xx)>=.4*px*px:
            values=rgb[y+yy,x+xx].astype(np.int16)
            # Between the crop detector's weak (60) and strong (100) colour
            # thresholds: pale blue ruling must not become an ownership anchor.
            pen[i] = np.median(np.ptp(values,axis=1))>80 or np.median(values.min(axis=1))<110
    support = credible & ink
    # The blank classifier deliberately ignores short isolated pen marks.
    # A two-dimensional letter/terminal stroke can still link to its page here.
    support |= pen[labels] & ink
    left_limit,right_limit = max(0,round(middle-12*px)),min(w,round(middle+12*px))
    left_counts = np.bincount(labels[:,:left_limit][support[:,:left_limit]],minlength=n)
    right_counts = np.bincount(labels[:,right_limit:][support[:,right_limit:]],minlength=n)
    anchors = np.c_[left_counts>=max(2,.2*px*px),right_counts>=max(2,.2*px*px)]
    backed = np.bincount(labels[support],minlength=n)>=max(2,.12*px*px)
    backed[0] = False
    near = max(1,round(2*px))
    roots = np.unique(labels[:,max(0,middle-near):middle+near])
    regions = ignored
    for root in roots:
        if root == 0:continue
        x0,y0,x1,y1 = boxes[root]
        if x1 < middle-near or x0 > middle+near:continue
        thin_unanchored = (not anchors[root].any() and
                          ((not pen[root] and y1-y0<px and x1-x0>4*(y1-y0)) or areas[root]<.4*px*px))
        if thin_unanchored:
            regions.append(dict(owner='artifact',box=boxes[root].tolist(),component_id=int(root),anchor_boxes=[None,None],anchor_distance_mm=[None,None]))
            continue
        distances = np.full(n,np.inf)
        distances[root]=0
        queue=[(0.,int(root))]
        reached=[np.inf,np.inf]
        evidence=[None,None]
        # Only supported writing may carry a connection onwards. Otherwise
        # rows of grid fragments could form a bridge all the way into an answer.
        allowed = backed.copy();allowed[root]=True
        while queue:
            cost,node=heapq.heappop(queue)
            if cost != distances[node] or cost>45*px:continue
            if anchors[node].any():
                for side in (0,1):
                    if anchors[node,side] and cost<reached[side]:
                        reached[side]=cost;evidence[side]=boxes[node].tolist()
                continue
            a,b,c,d = boxes[node]
            dx=np.maximum(0,np.maximum(boxes[:,0]-c,a-boxes[:,2]))
            dy=np.maximum(0,np.maximum(boxes[:,1]-d,b-boxes[:,3]))
            # Wide mathematical spaces are plausible between strong pen
            # regions; unsupported shadows get only the shorter word gap.
            gap = 32*px if pen[root] else 8*px
            neighbours=np.flatnonzero(allowed & (dx<=gap) & (dy<=2*px))
            for other in neighbours:
                step=float(np.hypot(dx[other],3*dy[other]))+1
                new=cost+step
                if new<distances[other]:
                    distances[other]=new;heapq.heappush(queue,(new,int(other)))
        if not np.isfinite(reached).any():owner='artifact'
        elif abs(reached[0]-reached[1])<=px:owner='uncertain'
        else:owner='left' if reached[0]<reached[1] else 'right'
        regions.append(dict(owner=owner,box=[int(x0),int(y0),int(x1),int(y1)],
            component_id=int(root),anchor_boxes=evidence,anchor_distance_mm=[round(c/px,2) if np.isfinite(c) else None for c in reached]))
    return regions,labels


def plan_masks(ink: np.ndarray, dpi: float, kind: str, *,
               content_ink: np.ndarray | None = None,
               rgb: np.ndarray | None = None,
               _remaining_content: dict | None = None) -> dict:
    """Produce rectangles in image coordinates and explicit review reasons."""
    h, w = ink.shape
    px = dpi / 25.4
    middle = w // 2
    reasons = []
    crossings = []
    left_end, right_start = middle, middle
    owned = (ink if content_ink is None else content_ink).copy()
    left_owned, right_owned = owned[:, :middle].copy(), owned[:, middle:].copy()
    if kind == 'a3':
        pad = round(2*px)
        crossings,group_labels = _fold_regions(ink,owned,dpi,rgb)
        for crossing in crossings:
            owner = crossing['owner']
            x0,y0,x1,y1 = crossing['box']
            if owner == 'artifact':
                continue
            if owner == 'left':
                left_end = max(left_end,min(w,x1+pad))
                if x1>middle:
                    right_owned[y0:y1,max(x0,middle)-middle:x1-middle] &= ~(ink[y0:y1,max(x0,middle):x1] & (group_labels[y0:y1,max(x0,middle):x1]==crossing['component_id']))
            elif owner == 'right':
                right_start = min(right_start,max(0,x0-pad))
                if x0<middle:
                    left_owned[y0:y1,x0:min(x1,middle)] &= ~(ink[y0:y1,x0:min(x1,middle)] & (group_labels[y0:y1,x0:min(x1,middle)]==crossing['component_id']))
            else:
                left_end = max(left_end,min(w,x1+pad))
                right_start = min(right_start,max(0,x0-pad))
                reasons.append('ambiguous-centre-writing')
        regions = [('left', [0,0,left_end,h], left_owned),
                   ('right', [right_start,0,w,h], right_owned)]
    else:
        regions = [('whole', [0,0,w,h], owned)]
        if kind == 'unknown':
            reasons.append('unknown-format-or-orientation')

    outputs = []
    if _remaining_content is not None:
        _remaining_content.update({side: content for side,_,content in regions})
    for side, rect, content in regions:
        state = 'content' if content.any() else 'blank'
        if kind == 'unknown':
            state = 'uncertain'
        outputs.append(dict(side=side, rect_px=rect, state=state,
                            keep=state != 'blank', ink_pixels=int(content.sum())))
    return dict(format=kind, size_px=[w,h], outputs=outputs, crossings=crossings,
                review_reasons=sorted(set(reasons)),
                extensions_mm=[round((left_end-middle)/px,2),
                               round((middle-right_start)/px,2)] if kind == 'a3' else [0,0])


def save_preview(rgb: np.ndarray, plan: dict, path: Path) -> None:
    im = Image.fromarray(rgb)
    im.thumbnail((1200, 1000))
    sx, sy = im.width/plan['size_px'][0], im.height/plan['size_px'][1]
    overlay = Image.new('RGBA', im.size)
    draw = ImageDraw.Draw(overlay)
    for output in plan['outputs']:
        x0,y0,x1,y1 = output['rect_px']
        box = (x0*sx,y0*sy,x1*sx-1,y1*sy-1)
        colour = '#187941' if output['side'] == 'left' else '#a832a6'
        draw.rectangle(box, outline=colour, width=3,
                       fill=(255,30,30,65) if not output['keep'] else None)
    for crossing in plan['crossings']:
        x0,y0,x1,y1 = crossing['box']
        draw.rectangle((x0*sx,y0*sy,x1*sx,y1*sy), outline='#999999' if crossing['owner']=='artifact' else '#e88b00', width=2)
    Image.alpha_composite(im.convert('RGBA'), overlay).convert('RGB').save(path, quality=90)


def process_pdf(job: tuple[str, str, int]) -> list[dict]:
    source_s, destination_s, dpi = job
    source, destination = Path(source_s), Path(destination_s)
    cv2.setNumThreads(1)
    digest = hashlib.sha256(source.read_bytes()).hexdigest()
    records = []
    with pymupdf.open(source) as doc, pymupdf.open() as result:
        for i, page in enumerate(doc):
            rotation = page.rotation
            kind = classify_format(page.rect.width, page.rect.height)
            from copienator.fast_page_split import analyse_pdf_page_fast
            plan,rgb = analyse_pdf_page_fast(page,dpi,kind)
            pixel_width,pixel_height = plan['size_px']
            plan.update(file=source.name, source_page=i+1, source_sha256=digest,
                        source_rotation=rotation, source_cropbox=list(page.cropbox))
            preview = f'{source.stem}-{i+1:03}.jpg'
            plan['preview'] = f'review/{preview}'
            save_preview(rgb, plan, destination/'review'/preview)
            # Normalize through a full-page PDF form: remove_rotation() loses
            # non-default CropBoxes on quarter-turn pages in some MuPDF builds.
            w,h = page.rect.width,page.rect.height
            normalized = pymupdf.open()
            flat = normalized.new_page(width=w,height=h)
            page.set_rotation(0)
            if page.get_contents():
                flat.show_pdf_page(flat.rect,doc,i,rotate=-rotation)
            for output in plan['outputs']:
                x0,y0,x1,y1 = output['rect_px']
                clip = pymupdf.Rect(x0/pixel_width*w,y0/pixel_height*h,
                                    x1/pixel_width*w,y1/pixel_height*h)
                output['crop_points'] = list(clip)
                if output['keep']:
                    target = result.new_page(width=clip.width,height=clip.height)
                    # Truly empty PDF pages have no content stream to show.
                    if flat.get_contents():
                        target.show_pdf_page(target.rect,normalized,0,clip=clip)
                    output['output_page'] = len(result)
            records.append(plan)
            normalized.close()
        # An entirely blank copy has no output PDF; its report is still kept.
        if len(result):
            result.save(destination/'pdfs'/source.name, garbage=3, deflate=True)
    if hashlib.sha256(source.read_bytes()).hexdigest() != digest:
        raise RuntimeError(f'Source changed during analysis: {source}')
    output_path = destination/'pdfs'/source.name
    expected = sum(o['keep'] for r in records for o in r['outputs'])
    if expected:
        with pymupdf.open(output_path) as check:
            if len(check) != expected or any(p.rect.is_empty for p in check):
                raise RuntimeError(f'Invalid output: {output_path}')
    return records


def run(source: Path, destination: Path, *, dpi: int = 200, workers: int = 1) -> dict:
    if dpi < 100 or workers < 1:
        raise ValueError('Use dpi >= 100 and workers >= 1')
    sources = sorted(source.glob('*.pdf')) if source.is_dir() else [source]
    if not sources or any(not p.is_file() or p.suffix.lower() != '.pdf' for p in sources):
        raise ValueError('No input PDFs found')
    # Refuse reuse: this makes accidental source replacement impossible.
    destination.mkdir(parents=True, exist_ok=False)
    (destination/'review').mkdir()
    (destination/'pdfs').mkdir()
    jobs = [(str(p),str(destination),dpi) for p in sources]
    records = []
    if workers == 1:
        batches = map(process_pdf, jobs)
        for batch in batches:
            records.extend(batch)
            print(f'{batch[0]["file"]}: {len(batch)} scanned pages', flush=True)
    else:
        with ProcessPoolExecutor(max_workers=workers,
                                 mp_context=multiprocessing.get_context('spawn')) as pool:
            for batch in pool.map(process_pdf,jobs):
                records.extend(batch)
                print(f'{batch[0]["file"]}: {len(batch)} scanned pages', flush=True)
    return write_report(records,destination,len(sources))


def write_report(records: list[dict], destination: Path, file_count: int) -> dict:
    """Rebuild the gallery from saved decisions without reanalysing scans."""
    summary = dict(files=file_count, source_pages=len(records),
        a3=sum(r['format']=='a3' for r in records), a4=sum(r['format']=='a4' for r in records),
        output_pages=sum(o['keep'] for r in records for o in r['outputs']),
        blank_removed=sum(not o['keep'] for r in records for o in r['outputs']),
        review_pages=sum(bool(r['review_reasons']) for r in records),
        extended_spreads=sum(any(r['extensions_mm']) for r in records))
    (destination/'report.json').write_text(json.dumps(dict(summary=summary,pages=records),indent=2)+'\n')
    cards = []
    for row in records:
        name = html.escape(row['file'])
        states = ', '.join(f'{o["side"]}: {o["state"]}' for o in row['outputs'])
        reasons = ', '.join(row['review_reasons']) or 'No review flag'
        links = ' · '.join(f'<a href="pdfs/{quote(row["file"])}#page={o["output_page"]}">{o["side"]} PDF</a>'
                           for o in row['outputs'] if o['keep'])
        cards.append(f'<article class="{"review" if row["review_reasons"] else "ok"}">'
            f'<h2>{name} / {row["source_page"]}</h2><p>{row["format"]} | {states}</p>'
            f'<p>Extensions L/R: {row["extensions_mm"]} mm | {reasons}</p><p>{links}</p>'
            f'<a href="{html.escape(row["preview"])}"><img loading="lazy" src="{html.escape(row["preview"])}"></a></article>')
    (destination/'index.html').write_text('<!doctype html><meta charset="utf-8">'
        '<title>Automatic page splitting - review</title><style>'
        'body{font:16px system-ui;margin:24px;background:#eee}main{display:grid;grid-template-columns:repeat(auto-fit,minmax(550px,1fr));gap:16px}'
        'article{background:white;padding:16px}article.review{border-top:5px solid #e88b00}img{width:100%}h2{font-size:18px}p{font-size:13px}</style>'
        '<h1>Automatic page splitting - experimental</h1><p>Source order; no booklet reordering. '
        'Green: left output. Purple: right/whole output. Orange: centre writing. Red shading: removed blank.</p>'
        f'<pre>{html.escape(json.dumps(summary,indent=2))}</pre><main>'+''.join(cards)+'</main>')
    return summary


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('input',type=Path)
    parser.add_argument('output',type=Path)
    parser.add_argument('--dpi',type=int,default=200)
    parser.add_argument('--workers',type=int,default=1)
    args = parser.parse_args(argv)
    print(json.dumps(run(args.input,args.output,dpi=args.dpi,workers=args.workers),indent=2))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())