"""Shared output helpers for image OCR results and batch manifests."""

from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any

from PIL import Image


def safe_stem(value: str) -> str:
    cleaned = re.sub(r"[^\w.-]+", "_", value, flags=re.UNICODE).strip("._")
    return cleaned or "result"


def atomic_write_text(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.tmp")
    temporary.write_text(content, encoding="utf-8")
    temporary.replace(path)


def atomic_write_json(path: Path, data: Any) -> None:
    atomic_write_text(path, json.dumps(data, ensure_ascii=False, indent=2))


def image_output_directory(
    output_root: Path,
    image_path: Path,
    *,
    batch_root: Path | None,
    recursive: bool,
) -> Path:
    base = output_root.expanduser().resolve() / "images"
    if batch_root is not None and recursive:
        relative_parent = image_path.parent.resolve().relative_to(batch_root.resolve())
        base /= relative_parent
    suffix = image_path.suffix.lower().lstrip(".") or "image"
    return base / safe_stem(f"{image_path.stem}_{suffix}")


def pdf_output_root(
    output_root: Path,
    pdf_path: Path,
    *,
    batch_root: Path | None,
    recursive: bool,
) -> Path:
    base = output_root.expanduser().resolve() / "pdfs"
    if batch_root is not None and recursive:
        relative_parent = pdf_path.parent.resolve().relative_to(batch_root.resolve())
        base /= relative_parent
    return base


def _save_asset(data: Any, path: Path) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    if isinstance(data, Image.Image):
        image = data
    else:
        import numpy as np

        array = np.asarray(data)
        if array.ndim not in (2, 3):
            raise TypeError(f"不支持的图片资源形状: {array.shape}")
        image = Image.fromarray(array.astype("uint8"))

    image_format = (path.suffix.lstrip(".") or "png").upper()
    if image_format == "JPG":
        image_format = "JPEG"
    if image_format not in {"PNG", "JPEG", "WEBP", "BMP", "TIFF"}:
        path = path.with_suffix(".png")
        image_format = "PNG"
    temporary = path.with_name(f".{path.name}.tmp")
    image.save(temporary, format=image_format)
    temporary.replace(path)
    return path


def save_image_ocr_outputs(
    result: Any,
    output_dir: Path,
    *,
    input_path: Path,
    benchmark: dict[str, Any],
) -> dict[str, str]:
    """Persist Markdown, plain text, Paddle JSON, and benchmark data."""
    output_dir.mkdir(parents=True, exist_ok=True)

    markdown_data = result.markdown
    if "res" in markdown_data and isinstance(markdown_data["res"], dict):
        markdown_data = markdown_data["res"]
    markdown_text = str(markdown_data.get("markdown_texts", "")).strip()
    asset_dir = output_dir / "assets"
    if asset_dir.exists():
        import shutil

        shutil.rmtree(asset_dir)
    for index, (original_path, image_data) in enumerate(
        (markdown_data.get("markdown_images") or {}).items(),
        start=1,
    ):
        original = str(original_path).replace("\\", "/")
        source_name = Path(original).name or f"image-{index:03d}.png"
        target = asset_dir / (
            f"{index:03d}-{safe_stem(Path(source_name).stem)}"
            f"{Path(source_name).suffix or '.png'}"
        )
        target = _save_asset(image_data, target)
        relative = Path(os.path.relpath(target, output_dir)).as_posix()
        markdown_text = markdown_text.replace(original, relative)
        markdown_text = markdown_text.replace(str(original_path), relative)

    plain_text = "\n\n".join(
        block.content.strip()
        for block in result["parsing_res_list"]
        if block.content.strip()
    )
    result_json = result.json
    payload = result_json.get("res", result_json) if isinstance(result_json, dict) else None
    if isinstance(payload, dict):
        payload["input_path"] = str(input_path)
        payload["source_type"] = "image_ocr"

    paths = {
        "output_dir": str(output_dir),
        "markdown": str(output_dir / "result.md"),
        "text": str(output_dir / "result.txt"),
        "json": str(output_dir / "result.json"),
        "benchmark": str(output_dir / "benchmark.json"),
    }
    atomic_write_text(Path(paths["markdown"]), markdown_text.rstrip() + "\n")
    atomic_write_text(Path(paths["text"]), plain_text.rstrip() + "\n")
    atomic_write_json(Path(paths["json"]), result_json)
    atomic_write_json(Path(paths["benchmark"]), benchmark)
    return paths