src/character_spike/assets.py

from __future__ import annotations

import json
import math
import random
import time
from collections import deque
from pathlib import Path
from typing import Any

from PIL import Image, ImageDraw, ImageFilter, ImageFont

from src.character_spike.schema import CANONICAL_STAGE_SIZE, EXPRESSIONS, default_character_package


def create_asset_package_from_probe_outputs(
    *,
    source_run_dir: str | Path,
    candidate_id: str,
    character_id: str,
    display_name: str,
    output_root: str | Path,
    seed: int = 42,
    remove_background: bool = True,
) -> dict[str, Any]:
    started = time.perf_counter()
    source = Path(source_run_dir)
    out_root = Path(output_root)
    run_dir = out_root / character_id
    character_dir = run_dir / "assets" / "characters" / character_id
    background_dir = run_dir / "assets" / "backgrounds"
    generated_dir = run_dir / "generated"
    character_dir.mkdir(parents=True, exist_ok=True)
    background_dir.mkdir(parents=True, exist_ok=True)
    generated_dir.mkdir(parents=True, exist_ok=True)

    package = default_character_package(character_id, display_name)
    package["metadata"]["source"] = "probe_asset_package"
    package["metadata"]["source_run_dir"] = str(source)
    package["metadata"]["candidate_id"] = candidate_id
    package_path = run_dir / "characters" / f"{character_id}.json"
    package_path.parent.mkdir(parents=True, exist_ok=True)
    package_path.write_text(json.dumps(package, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    assets: list[dict[str, Any]] = []
    missing_slots: list[str] = []
    for expression in EXPRESSIONS:
        source_path = source / "generated" / candidate_id / f"expression_{expression}" / "00.png"
        if not source_path.exists():
            source_path = source / "generated" / candidate_id / "expression_idle" / "00.png"
            missing_slots.append(expression)
        target_path = character_dir / f"{expression}.png"
        image = Image.open(source_path).convert("RGBA")
        processed = remove_flat_background(image) if remove_background else image
        normalized = normalize_to_stage_canvas(processed)
        normalized.save(target_path)
        assets.append(
            {
                "slot": expression,
                "path": str(target_path.relative_to(run_dir)),
                "source_path": str(source_path),
                "bytes": target_path.stat().st_size,
                "source": f"probe:{candidate_id}",
                "usable": expression not in missing_slots,
                "manual_score": None,
            }
        )

    background_path = background_dir / f"{character_id}_spike_background.png"
    _draw_mock_background(display_name, random.Random(seed)).save(background_path)
    grid_path = generated_dir / "asset_grid.png"
    make_thumbnail_grid([character_dir / f"{slot}.png" for slot in EXPRESSIONS], grid_path)

    manifest = {
        "schema_version": 1,
        "run_type": "probe_asset_package",
        "character_id": character_id,
        "display_name": display_name,
        "seed": seed,
        "created_at_unix": int(time.time()),
        "duration_seconds": round(time.perf_counter() - started, 3),
        "paths": {
            "run_dir": str(run_dir),
            "character_package": str(package_path.relative_to(run_dir)),
            "character_assets": str(character_dir.relative_to(run_dir)),
            "background": str(background_path.relative_to(run_dir)),
            "thumbnail_grid": str(grid_path.relative_to(run_dir)),
        },
        "assets": assets,
        "model_results": [],
        "qa": {
            "usable_assets": len([asset for asset in assets if asset["usable"]]),
            "total_assets": len(EXPRESSIONS),
            "needs_manual_review": missing_slots,
            "notes": [
                "Packaged from Modal probe outputs.",
                "Background removal uses a simple flat-background alpha pass; manual QA is still required."
                if remove_background
                else "Background was intentionally preserved as a no-matting fallback.",
            ],
        },
    }
    manifest_path = generated_dir / "manifest.json"
    manifest["paths"]["manifest"] = str(manifest_path.relative_to(run_dir))
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    write_report(manifest, generated_dir / "report.md")
    return manifest


def create_mock_asset_package(
    *,
    character_id: str,
    display_name: str,
    output_root: str | Path,
    seed: int = 42,
) -> dict[str, Any]:
    started = time.perf_counter()
    rng = random.Random(seed)
    out_root = Path(output_root)
    run_dir = out_root / character_id
    character_dir = run_dir / "assets" / "characters" / character_id
    background_dir = run_dir / "assets" / "backgrounds"
    generated_dir = run_dir / "generated"
    character_dir.mkdir(parents=True, exist_ok=True)
    background_dir.mkdir(parents=True, exist_ok=True)
    generated_dir.mkdir(parents=True, exist_ok=True)

    package = default_character_package(character_id, display_name)
    package["metadata"]["source"] = "mock_asset_package"
    package_path = run_dir / "characters" / f"{character_id}.json"
    package_path.parent.mkdir(parents=True, exist_ok=True)
    package_path.write_text(json.dumps(package, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    assets: list[dict[str, Any]] = []
    for index, expression in enumerate(EXPRESSIONS):
        path = character_dir / f"{expression}.png"
        image = _draw_mock_character(display_name, expression, index, rng)
        normalized = normalize_to_stage_canvas(image)
        normalized.save(path)
        assets.append(
            {
                "slot": expression,
                "path": str(path.relative_to(run_dir)),
                "bytes": path.stat().st_size,
                "source": "mock_pillow",
                "usable": True,
                "manual_score": None,
            }
        )

    background_path = background_dir / f"{character_id}_spike_background.png"
    _draw_mock_background(display_name, rng).save(background_path)
    grid_path = generated_dir / "asset_grid.png"
    make_thumbnail_grid([character_dir / f"{slot}.png" for slot in EXPRESSIONS], grid_path)

    manifest = {
        "schema_version": 1,
        "run_type": "mock_asset_package",
        "character_id": character_id,
        "display_name": display_name,
        "seed": seed,
        "created_at_unix": int(time.time()),
        "duration_seconds": round(time.perf_counter() - started, 3),
        "paths": {
            "run_dir": str(run_dir),
            "character_package": str(package_path.relative_to(run_dir)),
            "character_assets": str(character_dir.relative_to(run_dir)),
            "background": str(background_path.relative_to(run_dir)),
            "thumbnail_grid": str(grid_path.relative_to(run_dir)),
        },
        "assets": assets,
        "model_results": [],
        "qa": {
            "usable_assets": len(assets),
            "total_assets": len(EXPRESSIONS),
            "needs_manual_review": [],
            "notes": ["Mock assets validate packaging, postprocessing, manifest, and reporting without GPU cost."],
        },
    }
    manifest_path = generated_dir / "manifest.json"
    manifest["paths"]["manifest"] = str(manifest_path.relative_to(run_dir))
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    write_report(manifest, generated_dir / "report.md")
    return manifest


def normalize_to_stage_canvas(image: Image.Image, size: tuple[int, int] = CANONICAL_STAGE_SIZE) -> Image.Image:
    image = image.convert("RGBA")
    alpha = image.getchannel("A")
    bbox = alpha.getbbox()
    if bbox is None:
        return Image.new("RGBA", size, (0, 0, 0, 0))

    crop = image.crop(bbox)
    max_width = int(size[0] * 0.82)
    max_height = int(size[1] * 0.92)
    scale = min(max_width / crop.width, max_height / crop.height)
    new_size = (max(1, round(crop.width * scale)), max(1, round(crop.height * scale)))
    crop = crop.resize(new_size, Image.Resampling.LANCZOS)
    canvas = Image.new("RGBA", size, (0, 0, 0, 0))
    x = (size[0] - new_size[0]) // 2
    y = size[1] - new_size[1] - 24
    canvas.alpha_composite(crop, (x, y))
    return canvas


def remove_flat_background(
    image: Image.Image,
    *,
    tolerance: int = 42,
    feather_radius: float = 1.25,
) -> Image.Image:
    image = image.convert("RGBA")
    width, height = image.size
    sample_points = [
        (0, 0),
        (width - 1, 0),
        (0, height - 1),
        (width - 1, height - 1),
        (width // 2, 0),
        (width // 2, height - 1),
    ]
    pixels = image.load()
    samples = [pixels[x, y][:3] for x, y in sample_points]
    background = tuple(round(sum(channel) / len(samples)) for channel in zip(*samples))

    source_pixels = image.load()
    visited = bytearray(width * height)
    queue: deque[tuple[int, int]] = deque()

    def is_background(x: int, y: int) -> bool:
        rgb = source_pixels[x, y][:3]
        distance = max(abs(rgb[index] - background[index]) for index in range(3))
        return distance <= tolerance

    def push(x: int, y: int) -> None:
        index = y * width + x
        if visited[index] or not is_background(x, y):
            return
        visited[index] = 1
        queue.append((x, y))

    for x in range(width):
        push(x, 0)
        push(x, height - 1)
    for y in range(height):
        push(0, y)
        push(width - 1, y)

    while queue:
        x, y = queue.popleft()
        if x > 0:
            push(x - 1, y)
        if x < width - 1:
            push(x + 1, y)
        if y > 0:
            push(x, y - 1)
        if y < height - 1:
            push(x, y + 1)

    alpha = Image.new("L", image.size, 255)
    alpha_pixels = alpha.load()
    for y in range(height):
        row = y * width
        for x in range(width):
            if visited[row + x]:
                alpha_pixels[x, y] = 0
    if feather_radius > 0:
        alpha = alpha.filter(ImageFilter.GaussianBlur(feather_radius))
    image.putalpha(alpha)
    return image


def make_thumbnail_grid(paths: list[Path], output_path: str | Path, thumb_size: tuple[int, int] = (180, 240)) -> Path:
    output = Path(output_path)
    cols = 4
    rows = math.ceil(len(paths) / cols)
    grid = Image.new("RGB", (cols * thumb_size[0], rows * (thumb_size[1] + 26)), (17, 24, 39))
    draw = ImageDraw.Draw(grid)
    font = _font(16)
    for index, path in enumerate(paths):
        row, col = divmod(index, cols)
        x = col * thumb_size[0]
        y = row * (thumb_size[1] + 26)
        image = Image.open(path).convert("RGBA")
        image.thumbnail(thumb_size, Image.Resampling.LANCZOS)
        cell = Image.new("RGBA", thumb_size, (15, 23, 42, 255))
        px = (thumb_size[0] - image.width) // 2
        py = (thumb_size[1] - image.height) // 2
        cell.alpha_composite(image, (px, py))
        grid.paste(cell.convert("RGB"), (x, y))
        draw.text((x + 8, y + thumb_size[1] + 4), path.stem, fill=(226, 232, 240), font=font)
    output.parent.mkdir(parents=True, exist_ok=True)
    grid.save(output)
    return output


def write_report(manifest: dict[str, Any], path: str | Path) -> Path:
    output = Path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    report = render_report_markdown(manifest)
    output.write_text(report, encoding="utf-8")
    return output


def render_report_markdown(manifest: dict[str, Any]) -> str:
    qa = manifest.get("qa", {})
    model_results = manifest.get("model_results") or []
    assets = manifest.get("assets") or []
    lines = [
        f"# Character Generation Spike Report: {manifest.get('display_name', manifest.get('character_id', 'unknown'))}",
        "",
        "## Run",
        "",
        f"- Character ID: `{manifest.get('character_id')}`",
        f"- Run type: `{manifest.get('run_type')}`",
        f"- Duration: `{manifest.get('duration_seconds')}` seconds",
        f"- Seed: `{manifest.get('seed')}`",
        "",
        "## Asset QA",
        "",
        f"- Usable assets: `{qa.get('usable_assets', 0)}/{qa.get('total_assets', len(assets))}`",
        f"- Needs manual review: `{len(qa.get('needs_manual_review') or [])}`",
    ]
    for note in qa.get("notes") or []:
        lines.append(f"- Note: {note}")
    lines.extend(["", "## Model Results", ""])
    if model_results:
        lines.append("| Candidate | Mode | Cold/Warm | Images | Seconds | GPU | Status |")
        lines.append("| --- | --- | --- | ---: | ---: | --- | --- |")
        for result in model_results:
            lines.append(
                "| {candidate} | {mode} | {temperature} | {images} | {seconds} | {gpu} | {status} |".format(
                    candidate=result.get("candidate_id", ""),
                    mode=result.get("mode", ""),
                    temperature="cold" if result.get("loaded_before") is False else "warm",
                    images=result.get("image_count", 0),
                    seconds=result.get("duration_seconds", ""),
                    gpu=result.get("gpu", ""),
                    status=result.get("status", ""),
                )
            )
    else:
        lines.append("No remote model probes recorded yet.")

    lines.extend(["", "## Assets", ""])
    for asset in assets:
        lines.append(f"- `{asset.get('slot')}`: `{asset.get('path')}` ({asset.get('bytes')} bytes)")
    lines.append("")
    return "\n".join(lines)


def _draw_mock_character(display_name: str, expression: str, index: int, rng: random.Random) -> Image.Image:
    width, height = CANONICAL_STAGE_SIZE
    image = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    draw = ImageDraw.Draw(image)
    accent = _palette(index)
    skin = (255, 216, 204, 255)
    coat = (31 + rng.randrange(20), 41 + rng.randrange(20), 55 + rng.randrange(30), 255)
    hair = (190 + rng.randrange(45), 230 + rng.randrange(20), 245 + rng.randrange(10), 255)

    draw.ellipse((265, 970, 635, 1060), fill=(0, 0, 0, 72))
    draw.polygon([(250, 1080), (345, 660), (555, 660), (650, 1080)], fill=coat)
    draw.polygon([(345, 670), (450, 850), (555, 670), (595, 1080), (305, 1080)], fill=(12, 18, 31, 255))
    draw.rounded_rectangle((395, 600, 505, 730), radius=36, fill=(239, 184, 170, 255))
    draw.ellipse((255, 170, 645, 640), fill=skin)
    draw.ellipse((210, 95, 690, 520), fill=hair)
    draw.pieslice((210, 110, 690, 650), 185, 355, fill=(120, 190, 210, 255))
    draw.polygon([(300, 250), (370, 100), (420, 340)], fill=hair)
    draw.polygon([(430, 230), (510, 80), (540, 345)], fill=hair)
    draw.polygon([(535, 260), (610, 150), (625, 390)], fill=hair)
    draw.rounded_rectangle((385, 795, 515, 835), radius=18, fill=accent)

    _draw_expression(draw, expression, accent)
    font = _font(28)
    small = _font(22)
    draw.text((34, 34), display_name, fill=(238, 242, 255, 220), font=font)
    draw.text((34, 72), expression, fill=accent, font=small)
    return image


def _draw_expression(draw: ImageDraw.ImageDraw, expression: str, accent: tuple[int, int, int, int]) -> None:
    eye = (15, 23, 42, 255)
    if expression in {"smile", "happy"}:
        draw.arc((325, 385, 410, 445), 200, 340, fill=eye, width=8)
        draw.arc((490, 385, 575, 445), 200, 340, fill=eye, width=8)
        draw.arc((405, 505, 495, 570), 15, 165, fill=(159, 18, 57, 255), width=8)
    else:
        draw.ellipse((330, 380, 405, 455), fill=(248, 250, 252, 255))
        draw.ellipse((495, 380, 570, 455), fill=(248, 250, 252, 255))
        draw.ellipse((356, 398, 388, 438), fill=accent)
        draw.ellipse((521, 398, 553, 438), fill=accent)
        draw.ellipse((366, 410, 382, 433), fill=eye)
        draw.ellipse((531, 410, 547, 433), fill=eye)
        if expression == "worried":
            draw.arc((410, 525, 490, 580), 200, 340, fill=(159, 18, 57, 255), width=8)
        else:
            draw.arc((410, 505, 490, 552), 25, 155, fill=(159, 18, 57, 255), width=7)
    if expression == "thinking":
        draw.line((315, 350, 405, 360), fill=eye, width=7)
        draw.line((495, 360, 585, 350), fill=eye, width=7)
    elif expression == "worried":
        draw.line((315, 365, 405, 340), fill=eye, width=7)
        draw.line((495, 340, 585, 365), fill=eye, width=7)
    else:
        draw.line((320, 360, 405, 350), fill=eye, width=7)
        draw.line((495, 350, 580, 360), fill=eye, width=7)


def _draw_mock_background(display_name: str, rng: random.Random) -> Image.Image:
    width, height = 1600, 900
    image = Image.new("RGB", (width, height), (12, 18, 31))
    draw = ImageDraw.Draw(image)
    for y in range(height):
        shade = int(20 + 40 * y / height)
        draw.line((0, y, width, y), fill=(10, 16 + shade // 4, 26 + shade))
    for _ in range(70):
        x = rng.randrange(width)
        y = rng.randrange(height)
        radius = rng.randrange(1, 4)
        draw.ellipse((x, y, x + radius, y + radius), fill=(148, 163, 184))
    for x in range(0, width, 120):
        draw.line((x, 0, x + 320, height), fill=(255, 255, 255, 16), width=1)
    font = _font(44)
    draw.text((56, 56), f"{display_name} / spike background", fill=(226, 232, 240), font=font)
    return image


def _palette(index: int) -> tuple[int, int, int, int]:
    colors = [
        (103, 232, 249, 255),
        (125, 211, 252, 255),
        (250, 204, 21, 255),
        (244, 114, 182, 255),
        (167, 243, 208, 255),
        (196, 181, 253, 255),
        (251, 146, 60, 255),
        (129, 140, 248, 255),
    ]
    return colors[index % len(colors)]


def _font(size: int) -> ImageFont.ImageFont:
    try:
        return ImageFont.truetype("arial.ttf", size)
    except OSError:
        return ImageFont.load_default()