"""Unified suffix-based routing for PP-OCRv6 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 .result_adapter import adapt_ocr_result
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"
    if path.suffix.lower() in IMAGE_EXTENSIONS:
        return "image"
    if path.suffix.lower() 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 _predict_kwargs(args) -> dict[str, Any]:
    return {
        key: value
        for key, value in {
            "text_det_limit_side_len": args.text_det_limit_side_len,
            "text_det_limit_type": args.text_det_limit_type,
            "text_det_thresh": args.text_det_thresh,
            "text_det_box_thresh": args.text_det_box_thresh,
            "text_det_unclip_ratio": args.text_det_unclip_ratio,
            "text_rec_score_thresh": args.text_rec_score_thresh,
            "return_word_box": args.return_word_box,
        }.items()
        if value is not None
    }


def process_image_file(
    path: Path,
    *,
    args,
    provider: PipelineProvider,
    logger: logging.Logger,
    project_root: Path,
    run_warmup: bool,
    batch_root: Path | None,
) -> FileProcessResult:
    del project_root
    file_started = time.perf_counter()
    logger.info("FILE_ROUTED path=%s kind=image", path)
    pipeline = provider.get()
    predict_kwargs = _predict_kwargs(args)

    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()
        results = pipeline.predict(str(path), **predict_kwargs)
        provider.synchronize()
        elapsed = time.perf_counter() - started
        inference_times.append(elapsed)
        if not results:
            raise RuntimeError("PP-OCRv6 未返回图片结果")
        result = results[0]
        logger.info("INFERENCE_COMPLETED path=%s round=%d/%d seconds=%.3f", path, index + 1, args.rounds, elapsed)
    assert result is not None

    adapt_started = time.perf_counter()
    normalized = adapt_ocr_result(
        result,
        input_path=path,
        source_type="image_ocr",
        language=provider.config.lang,
        detection_model=provider.config.text_detection_model_name,
        recognition_model=provider.config.text_recognition_model_name,
        model_size=provider.config.model_size,
    )
    adapt_seconds = time.perf_counter() - adapt_started
    summary = normalized["summary"]
    benchmark = {
        "timestamp": datetime.now().astimezone().isoformat(),
        **provider.metadata(),
        "image_path": str(path),
        "image": normalized["image"],
        "ocr_summary": 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),
        },
        "result_adapt_seconds": round(adapt_seconds, 3),
        "gpu_memory": provider.gpu_memory(),
        "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,
        normalized,
        output_dir,
        input_path=path,
        benchmark=benchmark,
        save_raw_result=args.save_raw_result,
        save_visualization=args.save_visualization,
    )
    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 = args.benchmark_json.expanduser().resolve()
        atomic_write_json(explicit, benchmark)
        output_paths["explicit_benchmark"] = str(explicit)

    logger.info(
        "IMAGE_COMPLETED path=%s detected_lines=%d non_empty_lines=%d mean_score=%s inference_mean_seconds=%.3f file_total_seconds=%.3f output=%s",
        path,
        summary["detected_lines"],
        summary["non_empty_lines"],
        summary["mean_score"],
        statistics.fmean(inference_times),
        total_seconds,
        output_dir,
    )
    if not args.no_result:
        for line in normalized["lines"]:
            logger.info("OCR_LINE path=%s index=%d score=%s text=%s", path, line["index"], line["score"], line["text"].replace("\n", "\\n"))
    return FileProcessResult(path, "image", "completed", total_seconds, {**summary, **output_paths})


def process_pdf_file(path: Path, *, args, provider: PipelineProvider, logger: logging.Logger, batch_root: Path | None) -> FileProcessResult:
    started = time.perf_counter()
    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()
    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", path, preflight["page_count"], len(preflight["selected_pages"]))
    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=_predict_kwargs(args),
        logger=logger,
    )
    total = time.perf_counter() - started
    logger.info("PDF_COMPLETED path=%s status=%s text_pages=%d ocr_pages=%d failed_pages=%s total_seconds=%.3f", path, summary["status"], summary["text_pages"], summary["ocr_pages"], summary["failed_pages"], total)
    return FileProcessResult(path, "pdf", summary["status"], total, summary, 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)
    raise ValueError(f"不支持的文件类型: {path.suffix or '<无后缀>'}")


def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
    program_started = time.perf_counter()
    input_path = args.input.expanduser().resolve()
    if detect_input_kind(input_path) != "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", 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 total_seconds=%.3f", result.path, result.kind, result.status, 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", input_path)
        return 1
    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:
            return 130
        except Exception as exc:
            failures.append({"path": str(path), "error": f"{type(exc).__name__}: {exc}"})
            logger.exception("FILE_FAILED path=%s", path)
            if args.fail_fast:
                break
    program_total = time.perf_counter() - program_started
    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"
    atomic_write_json(
        manifest_path,
        {
            "input_directory": str(input_path),
            "recursive": args.recursive,
            "device": provider.resolved_device,
            "model_config": provider.model_config(),
            "discovered_files": len(files),
            "completed_files": len(results),
            "failed_files": len(failures),
            "program_total_seconds": round(program_total, 3),
            "results": [{"path": str(item.path), "kind": item.kind, "status": item.status, "seconds": round(item.seconds, 3), "outputs": item.details} for item in results],
            "failures": failures,
        },
    )
    logger.info("DIRECTORY_SUMMARY discovered=%d completed=%d failed=%d manifest=%s", len(files), len(results), len(failures), manifest_path)
    return 0 if not failures else 3


def run_verify(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
    del args, project_root
    started = time.perf_counter()
    try:
        provider.prepare()
        paddle = provider._paddle
        if provider.config.device == "gpu":
            result = paddle.matmul(paddle.ones([1024, 1024]), paddle.ones([1024, 1024]))
            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