"""Unified suffix-based routing for files and directories."""

from __future__ import annotations

import logging
import statistics
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

from .output import (
    atomic_write_json,
    image_output_directory,
    pdf_output_root,
    safe_stem,
    save_image_ocr_outputs,
)
from .pdf import preflight_pdf, process_pdf
from .pdf_text import TextLayerPolicy
from .runtime import PipelineProvider

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".webp"}
PDF_EXTENSIONS = {".pdf"}
SUPPORTED_EXTENSIONS = IMAGE_EXTENSIONS | PDF_EXTENSIONS


@dataclass
class FileProcessResult:
    path: Path
    kind: str
    status: str
    seconds: float
    details: dict[str, Any]
    exit_code: int = 0


def detect_input_kind(path: Path) -> str:
    if path.is_dir():
        return "directory"
    suffix = path.suffix.lower()
    if suffix in IMAGE_EXTENSIONS:
        return "image"
    if suffix in PDF_EXTENSIONS:
        return "pdf"
    return "unsupported"


def discover_supported_files(directory: Path, *, recursive: bool, output_root: Path) -> list[Path]:
    iterator = directory.rglob("*") if recursive else directory.glob("*")
    output_root = output_root.expanduser().resolve()
    files: list[Path] = []
    for path in iterator:
        if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS:
            continue
        resolved = path.resolve()
        try:
            resolved.relative_to(output_root)
        except ValueError:
            files.append(resolved)
    return sorted(files, key=lambda value: str(value).casefold())


def _result_summary(result: list[Any]) -> dict[str, Any]:
    first = result[0]
    blocks = first["parsing_res_list"]
    return {
        "width": first.get("width"),
        "height": first.get("height"),
        "layout_boxes": len(first.get("layout_det_res", {}).get("boxes", [])),
        "parsed_blocks": len(blocks),
        "non_empty_blocks": sum(bool(block.content.strip()) for block in blocks),
    }


def process_image_file(
    path: Path,
    *,
    args,
    provider: PipelineProvider,
    logger: logging.Logger,
    project_root: Path,
    run_warmup: bool,
    batch_root: Path | None,
) -> FileProcessResult:
    file_started = time.perf_counter()
    logger.info("FILE_ROUTED path=%s kind=image", path)
    pipeline = provider.get()
    predict_kwargs = {
        key: value
        for key, value in {
            "max_new_tokens": getattr(args, "max_new_tokens", None),
            "min_pixels": getattr(args, "min_pixels", None),
            "max_pixels": getattr(args, "max_pixels", None),
        }.items()
        if value is not None
    }

    warmup_times: list[float] = []
    if run_warmup:
        for index in range(args.warmup):
            started = time.perf_counter()
            pipeline.predict(str(path), **predict_kwargs)
            provider.synchronize()
            elapsed = time.perf_counter() - started
            warmup_times.append(elapsed)
            logger.info(
                "WARMUP_COMPLETED path=%s round=%d/%d seconds=%.3f",
                path,
                index + 1,
                args.warmup,
                elapsed,
            )

    inference_times: list[float] = []
    result = None
    for index in range(args.rounds):
        provider.synchronize()
        started = time.perf_counter()
        result = pipeline.predict(str(path), **predict_kwargs)
        provider.synchronize()
        elapsed = time.perf_counter() - started
        inference_times.append(elapsed)
        logger.info(
            "INFERENCE_COMPLETED path=%s round=%d/%d seconds=%.3f",
            path,
            index + 1,
            args.rounds,
            elapsed,
        )

    if not result:
        raise RuntimeError("OCR pipeline 未返回图片结果")

    summary = _result_summary(result)
    processing_seconds = time.perf_counter() - file_started
    benchmark = {
        "timestamp": datetime.now().astimezone().isoformat(),
        **provider.metadata(),
        "image_path": str(path),
        "image": summary,
        "warmup_seconds": [round(value, 3) for value in warmup_times],
        "inference_seconds": {
            "all": [round(value, 3) for value in inference_times],
            "min": round(min(inference_times), 3),
            "max": round(max(inference_times), 3),
            "mean": round(statistics.fmean(inference_times), 3),
            "median": round(statistics.median(inference_times), 3),
            "stdev": round(statistics.pstdev(inference_times), 3),
        },
        "gpu_memory": provider.gpu_memory(),
        "processing_seconds": round(processing_seconds, 3),
        "export_seconds": 0.0,
        "file_total_seconds": 0.0,
    }
    output_dir = image_output_directory(
        args.output,
        path,
        batch_root=batch_root,
        recursive=args.recursive,
    )
    export_started = time.perf_counter()
    output_paths = save_image_ocr_outputs(
        result[0],
        output_dir,
        input_path=path,
        benchmark=benchmark,
    )
    export_seconds = time.perf_counter() - export_started
    total_seconds = time.perf_counter() - file_started
    benchmark["export_seconds"] = round(export_seconds, 3)
    benchmark["file_total_seconds"] = round(total_seconds, 3)
    atomic_write_json(Path(output_paths["benchmark"]), benchmark)
    if batch_root is None and args.benchmark_json:
        explicit_benchmark = args.benchmark_json.expanduser().resolve()
        atomic_write_json(explicit_benchmark, benchmark)
        output_paths["explicit_benchmark"] = str(explicit_benchmark)

    logger.info(
        "IMAGE_COMPLETED path=%s width=%s height=%s layout_boxes=%d parsed_blocks=%d inference_mean_seconds=%.3f export_seconds=%.3f file_total_seconds=%.3f output=%s benchmark=%s",
        path,
        summary["width"],
        summary["height"],
        summary["layout_boxes"],
        summary["parsed_blocks"],
        statistics.fmean(inference_times),
        export_seconds,
        total_seconds,
        output_paths["output_dir"],
        output_paths["benchmark"],
    )
    if not args.no_result:
        for index, block in enumerate(result[0]["parsing_res_list"], 1):
            logger.info(
                "OCR_BLOCK path=%s index=%d label=%s bbox=%s content=%s",
                path,
                index,
                block.label,
                block.bbox,
                block.content.replace("\r", "").replace("\n", "\\n"),
            )

    return FileProcessResult(
        path=path,
        kind="image",
        status="completed",
        seconds=total_seconds,
        details={**summary, **output_paths},
    )


def process_pdf_file(
    path: Path,
    *,
    args,
    provider: PipelineProvider,
    logger: logging.Logger,
    batch_root: Path | None,
) -> FileProcessResult:
    file_started = time.perf_counter()
    logger.info("FILE_ROUTED path=%s kind=pdf mode=%s", path, args.pdf_mode)
    output_root = pdf_output_root(
        args.output,
        path,
        batch_root=batch_root,
        recursive=args.recursive,
    )
    manifest_exists = (output_root / safe_stem(path.stem) / "manifest.json").is_file()
    # Directory jobs auto-resume existing PDF manifests so rerunning a batch is
    # safe. Single-file jobs still require an explicit --resume.
    resume = args.resume if batch_root is None else manifest_exists and not args.overwrite
    preflight = preflight_pdf(
        pdf_path=path,
        output_root=output_root,
        pages=args.pages,
        dpi=args.dpi,
        password=args.password,
        resume=resume,
        overwrite=args.overwrite,
    )
    logger.info(
        "PDF_PREFLIGHT_COMPLETED path=%s page_count=%d selected_pages=%d output=%s mode=%s",
        path,
        preflight["page_count"],
        len(preflight["selected_pages"]),
        preflight["document_dir"],
        args.pdf_mode,
    )
    policy = TextLayerPolicy(
        min_chars=args.text_min_chars,
        min_printable_ratio=args.text_min_printable_ratio,
        min_content_ratio=args.text_min_content_ratio,
        max_replacement_ratio=args.text_max_replacement_ratio,
        min_chars_per_megapixel=args.text_min_density,
    )
    if args.pdf_mode == "ocr":
        provider.prepare()
    summary = process_pdf(
        provider=provider,
        pdf_path=path,
        output_root=output_root,
        mode=args.pdf_mode,
        text_policy=policy,
        pages=args.pages,
        dpi=args.dpi,
        password=args.password,
        resume=resume,
        overwrite=args.overwrite,
        keep_rendered=args.keep_rendered,
        fail_fast=args.fail_fast,
        predict_kwargs={
            key: value
            for key, value in {
                "max_new_tokens": args.max_new_tokens,
                "min_pixels": args.min_pixels,
                "max_pixels": args.max_pixels,
            }.items()
            if value is not None
        },
        logger=logger,
    )
    total_seconds = time.perf_counter() - file_started
    logger.info(
        "PDF_COMPLETED path=%s status=%s text_pages=%d ocr_pages=%d failed_pages=%s model_used=%s model_initialized_during_task=%s resume=%s file_total_seconds=%.3f output=%s",
        path,
        summary["status"],
        summary["text_pages"],
        summary["ocr_pages"],
        summary["failed_pages"],
        summary["model_used"],
        summary["model_initialized_during_task"],
        resume,
        total_seconds,
        summary["document_dir"],
    )
    return FileProcessResult(
        path=path,
        kind="pdf",
        status=summary["status"],
        seconds=total_seconds,
        details=summary,
        exit_code=0 if not summary["failed_pages"] else 3,
    )


def process_single_file(
    path: Path,
    *,
    args,
    provider: PipelineProvider,
    logger: logging.Logger,
    project_root: Path,
    run_image_warmup: bool,
    batch_root: Path | None = None,
) -> FileProcessResult:
    path = path.expanduser().resolve()
    if not path.is_file():
        raise FileNotFoundError(f"文件不存在: {path}")
    kind = detect_input_kind(path)
    if kind == "image":
        return process_image_file(
            path,
            args=args,
            provider=provider,
            logger=logger,
            project_root=project_root,
            run_warmup=run_image_warmup,
            batch_root=batch_root,
        )
    if kind == "pdf":
        return process_pdf_file(
            path,
            args=args,
            provider=provider,
            logger=logger,
            batch_root=batch_root,
        )
    supported = ", ".join(sorted(SUPPORTED_EXTENSIONS))
    raise ValueError(f"不支持的文件类型: {path.suffix or '<无后缀>'}；支持: {supported}")


def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
    program_started = time.perf_counter()
    input_path = args.input.expanduser().resolve()
    kind = detect_input_kind(input_path)

    if kind != "directory":
        try:
            result = process_single_file(
                input_path,
                args=args,
                provider=provider,
                logger=logger,
                project_root=project_root,
                run_image_warmup=True,
            )
        except KeyboardInterrupt:
            logger.warning("PROGRAM_INTERRUPTED input=%s resume_hint=--resume", input_path)
            return 130
        except Exception as exc:
            logger.exception("FILE_FAILED path=%s error=%s", input_path, exc)
            return 1
        logger.info(
            "PROGRAM_COMPLETED input=%s kind=%s status=%s file_seconds=%.3f program_total_seconds=%.3f",
            result.path,
            result.kind,
            result.status,
            result.seconds,
            time.perf_counter() - program_started,
        )
        return result.exit_code

    files = discover_supported_files(
        input_path,
        recursive=args.recursive,
        output_root=args.output,
    )
    if not files:
        logger.error("NO_SUPPORTED_FILES directory=%s recursive=%s", input_path, args.recursive)
        return 1
    logger.info(
        "DIRECTORY_PLAN directory=%s recursive=%s files=%d image_files=%d pdf_files=%d",
        input_path,
        args.recursive,
        len(files),
        sum(path.suffix.lower() in IMAGE_EXTENSIONS for path in files),
        sum(path.suffix.lower() in PDF_EXTENSIONS for path in files),
    )

    results: list[FileProcessResult] = []
    failures: list[dict[str, str]] = []
    image_warmup_pending = True
    for index, path in enumerate(files, 1):
        logger.info("DIRECTORY_PROGRESS_START progress=%d/%d path=%s", index, len(files), path)
        try:
            result = process_single_file(
                path,
                args=args,
                provider=provider,
                logger=logger,
                project_root=project_root,
                run_image_warmup=image_warmup_pending,
                batch_root=input_path,
            )
            results.append(result)
            if result.kind == "image":
                image_warmup_pending = False
            if result.exit_code:
                failures.append({"path": str(path), "error": result.status})
        except KeyboardInterrupt:
            logger.warning("PROGRAM_INTERRUPTED path=%s progress=%d/%d", path, index, len(files))
            return 130
        except Exception as exc:
            failures.append({"path": str(path), "error": f"{type(exc).__name__}: {exc}"})
            logger.exception("FILE_FAILED path=%s progress=%d/%d", path, index, len(files))
            if args.fail_fast:
                break
        logger.info("DIRECTORY_PROGRESS_END progress=%d/%d path=%s", index, len(files), path)

    image_results = [result for result in results if result.kind == "image"]
    pdf_results = [result for result in results if result.kind == "pdf"]
    total_file_seconds = sum(result.seconds for result in results)
    program_total = time.perf_counter() - program_started
    batch_manifest_path = (
        args.output.expanduser().resolve()
        / "batches"
        / f"{safe_stem(input_path.name or 'batch')}-{datetime.now():%Y%m%d-%H%M%S-%f}.json"
    )
    batch_manifest = {
        "input_directory": str(input_path),
        "recursive": args.recursive,
        "device": provider.resolved_device,
        "discovered_files": len(files),
        "completed_files": len(results),
        "failed_files": len(failures),
        "image_files": len(image_results),
        "pdf_files": len(pdf_results),
        "model_init_seconds": round(provider.model_init_seconds, 3),
        "total_file_seconds": round(total_file_seconds, 3),
        "program_total_seconds": round(program_total, 3),
        "results": [
            {
                "path": str(result.path),
                "kind": result.kind,
                "status": result.status,
                "seconds": round(result.seconds, 3),
                "exit_code": result.exit_code,
                "outputs": result.details,
            }
            for result in results
        ],
        "failures": failures,
    }
    atomic_write_json(batch_manifest_path, batch_manifest)
    logger.info(
        "DIRECTORY_SUMMARY discovered=%d completed=%d failed=%d images_completed=%d pdfs_completed=%d total_file_seconds=%.3f program_total_seconds=%.3f model_init_seconds=%.3f manifest=%s",
        len(files),
        len(results),
        len(failures),
        len(image_results),
        len(pdf_results),
        total_file_seconds,
        program_total,
        provider.model_init_seconds,
        batch_manifest_path,
    )
    return 0 if not failures else 3


def run_verify(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
    started = time.perf_counter()
    try:
        provider.prepare()
        paddle = provider._paddle
        if provider.config.device == "gpu":
            tensor = paddle.ones([1024, 1024], dtype="float32")
            result = paddle.matmul(tensor, tensor)
            provider.synchronize()
            logger.info("GPU_SMOKE_TEST shape=%s", list(result.shape))
        else:
            from paddle import core

            logger.info(
                "CPU_SMOKE_TEST onednn=%s mkldnn=%s",
                core.is_compiled_with_onednn(),
                core.is_compiled_with_mkldnn(),
            )
        logger.info(
            "VERIFY_COMPLETED metadata=%s seconds=%.3f",
            provider.metadata(),
            time.perf_counter() - started,
        )
        return 0
    except Exception:
        logger.exception("VERIFY_FAILED")
        return 1