from __future__ import annotations

import os
import logging
import importlib.util
from contextlib import contextmanager
from dataclasses import asdict, dataclass
from typing import Any, Dict, List, Optional, Tuple


@dataclass
class LayoutBlock:
    block_id: str
    page: int
    order: int
    text: str
    block_type: str = "text"
    label_name: str = "TEXT"
    source_ref: Optional[str] = None
    provenance: Optional[Dict[str, Any]] = None


_CONVERTER_CACHE: Dict[Tuple[Any, ...], Any] = {}
_DOCLING_CAPTURED_LOGGERS: Tuple[str, ...] = ("docling.pipeline.standard_pdf_pipeline",)


@dataclass
class DoclingParseOptions:
    profile: str = "balanced"
    do_ocr: bool = True
    force_full_page_ocr: bool = True
    backend: str = "torch"
    do_table_structure: bool = True
    do_picture_description: bool = False
    enable_remote_services: bool = False
    image_mode: str = "placeholder"
    quiet: bool = True


class _DoclingLogCaptureHandler(logging.Handler):
    def __init__(self) -> None:
        super().__init__(level=logging.WARNING)
        self.records: List[logging.LogRecord] = []

    def emit(self, record: logging.LogRecord) -> None:
        self.records.append(record)


DOCLING_PROFILE_DEFAULTS: Dict[str, Dict[str, Any]] = {
    "balanced": {
        "do_ocr": True,
        "force_full_page_ocr": True,
        "backend": "torch",
        "do_table_structure": True,
        "do_picture_description": False,
        "enable_remote_services": False,
        "image_mode": "placeholder",
        "quiet": True,
    },
    "image_manual": {
        "do_ocr": True,
        "force_full_page_ocr": True,
        "backend": "torch",
        "do_table_structure": True,
        "do_picture_description": True,
        "enable_remote_services": False,
        "image_mode": "placeholder",
        "quiet": True,
    },
    "fast_text": {
        "do_ocr": False,
        "force_full_page_ocr": False,
        "backend": "onnxruntime",
        "do_table_structure": True,
        "do_picture_description": False,
        "enable_remote_services": False,
        "image_mode": "placeholder",
        "quiet": True,
    },
}


def _env_bool(name: str, default: bool) -> bool:
    fallback = "1" if default else "0"
    raw = os.getenv(name, fallback).strip().lower()
    return raw in {"1", "true", "yes", "on"}


def _resolve_docling_parse_options(docling_options: Optional[Dict[str, Any]] = None) -> DoclingParseOptions:
    user = dict(docling_options or {})
    profile_name = str(user.pop("profile", os.getenv("NAVEXA_DOCLING_PROFILE", "balanced"))).strip().lower() or "balanced"
    profile = DOCLING_PROFILE_DEFAULTS.get(profile_name, DOCLING_PROFILE_DEFAULTS["balanced"])

    settings: Dict[str, Any] = dict(profile)
    settings.update(
        {
            "do_ocr": _env_bool("NAVEXA_DOCLING_OCR", bool(profile.get("do_ocr", True))),
            "force_full_page_ocr": _env_bool(
                "NAVEXA_DOCLING_FORCE_FULL_PAGE_OCR",
                bool(profile.get("force_full_page_ocr", True)),
            ),
            "backend": os.getenv("NAVEXA_RAPIDOCR_BACKEND", str(profile.get("backend", "torch"))).strip().lower() or "torch",
            "do_table_structure": _env_bool(
                "NAVEXA_DOCLING_TABLE_STRUCTURE",
                bool(profile.get("do_table_structure", True)),
            ),
            "do_picture_description": _env_bool(
                "NAVEXA_DOCLING_PICTURE_DESCRIPTION",
                bool(profile.get("do_picture_description", False)),
            ),
            "enable_remote_services": _env_bool(
                "NAVEXA_DOCLING_REMOTE_SERVICES",
                bool(profile.get("enable_remote_services", False)),
            ),
            "image_mode": os.getenv("NAVEXA_DOCLING_IMAGE_MODE", str(profile.get("image_mode", "placeholder"))).strip().lower()
            or "placeholder",
            "quiet": _env_bool("NAVEXA_DOCLING_QUIET", bool(profile.get("quiet", True))),
        }
    )
    settings.update(user)

    image_mode = str(settings.get("image_mode", "placeholder")).strip().lower() or "placeholder"
    if image_mode not in {"placeholder", "embedded", "referenced", "none"}:
        image_mode = "placeholder"
    settings["image_mode"] = image_mode
    settings["backend"] = str(settings.get("backend", "torch")).strip().lower() or "torch"
    if settings["backend"] not in {"torch", "onnxruntime"}:
        settings["backend"] = "torch"

    return DoclingParseOptions(
        profile=profile_name if profile_name in DOCLING_PROFILE_DEFAULTS else "balanced",
        do_ocr=bool(settings.get("do_ocr", True)),
        force_full_page_ocr=bool(settings.get("force_full_page_ocr", True)),
        backend=settings["backend"],
        do_table_structure=bool(settings.get("do_table_structure", True)),
        do_picture_description=bool(settings.get("do_picture_description", False)),
        enable_remote_services=bool(settings.get("enable_remote_services", False)),
        image_mode=settings["image_mode"],
        quiet=bool(settings.get("quiet", True)),
    )


def resolve_docling_options_for_pipeline(docling_options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    return asdict(_resolve_docling_parse_options(docling_options))


def _bbox_to_json(bbox: Any) -> Any:
    if bbox is None:
        return None
    if hasattr(bbox, "model_dump"):
        try:
            return bbox.model_dump()
        except Exception:
            pass
    if hasattr(bbox, "__dict__"):
        try:
            return dict(bbox.__dict__)
        except Exception:
            pass
    if isinstance(bbox, (list, tuple)):
        return list(bbox)
    return str(bbox)


def _build_provenance_summary(prov: List[Any], fallback_page: int) -> Dict[str, Any]:
    pages: List[int] = []
    for p in prov:
        page_no = getattr(p, "page_no", None)
        if isinstance(page_no, int):
            pages.append(int(page_no))

    if pages:
        page_start = min(pages)
        page_end = max(pages)
    else:
        page_start = fallback_page
        page_end = fallback_page

    bbox = _bbox_to_json(getattr(prov[0], "bbox", None)) if prov else None
    return {"page_start": page_start, "page_end": page_end, "bbox": bbox}


def _synthesize_blocks_from_text_map(page_texts: Dict[int, str]) -> List[LayoutBlock]:
    synthetic: List[LayoutBlock] = []
    for page in sorted(page_texts):
        txt = (page_texts.get(page) or "").strip()
        if not txt:
            continue
        synthetic.append(
            LayoutBlock(
                block_id=f"s{page}_0",
                page=int(page),
                order=0,
                text=txt,
                block_type="synthetic",
                label_name="SYNTHETIC_TEXT",
                source_ref=None,
                provenance={"page_start": int(page), "page_end": int(page), "bbox": None},
            )
        )
    return synthetic


def _fallback_text_from_markdown(markdown: str) -> str:
    lines = []
    for line in markdown.splitlines():
        raw = line.strip()
        if not raw:
            continue
        if raw == "<!-- image -->":
            continue
        lines.append(line.rstrip())
    return "\n".join(lines).strip()


def _safe_env() -> None:
    os.environ["MKL_THREADING_LAYER"] = "GNU"
    os.environ["OMP_NUM_THREADS"] = "1"
    os.environ["MKL_SERVICE_FORCE_INTEL"] = "1"
    os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def _configure_docling_logging(*, quiet: bool) -> None:
    if not quiet:
        return
    for logger_name in ("RapidOCR", "rapidocr"):
        target = logging.getLogger(logger_name)
        target.setLevel(logging.WARNING)
        for handler in target.handlers:
            handler.setLevel(logging.WARNING)


@contextmanager
def _capture_docling_pipeline_logs() -> List[logging.LogRecord]:
    handler = _DoclingLogCaptureHandler()
    logger_states: List[Tuple[logging.Logger, List[logging.Handler], int, bool]] = []
    try:
        for logger_name in _DOCLING_CAPTURED_LOGGERS:
            target = logging.getLogger(logger_name)
            logger_states.append((target, list(target.handlers), target.level, target.propagate))
            target.handlers = [handler]
            target.setLevel(logging.WARNING)
            target.propagate = False
        yield handler.records
    finally:
        for target, old_handlers, old_level, old_propagate in logger_states:
            target.handlers = old_handlers
            target.setLevel(old_level)
            target.propagate = old_propagate


def _normalize_docling_status(status: Any) -> Optional[str]:
    if status is None:
        return None
    value = getattr(status, "value", None)
    if isinstance(value, str) and value.strip():
        return value.strip().lower()
    text = str(status).strip().lower()
    if not text:
        return None
    return text.split(".")[-1]


def _normalize_docling_message(message: Any) -> str:
    return " ".join(str(message or "").split()).strip()


def _dedupe_messages(messages: List[str]) -> List[str]:
    seen = set()
    unique: List[str] = []
    for message in messages:
        if not message or message in seen:
            continue
        seen.add(message)
        unique.append(message)
    return unique


def _infer_docling_issue_stage(details: List[str]) -> str:
    for detail in details:
        lowered = detail.lower()
        if lowered.startswith("stage "):
            parts = detail.split()
            if len(parts) >= 2:
                return parts[1].strip().lower()
        if "ocr" in lowered:
            return "ocr"
        if "picture" in lowered:
            return "picture_description"
        if "table" in lowered:
            return "table_structure"
        if "backend" in lowered:
            return "backend"
    return "docling"


def _build_docling_report(result: Any, captured_records: List[logging.LogRecord]) -> Dict[str, Any]:
    details: List[str] = []
    for record in captured_records:
        message = _normalize_docling_message(record.getMessage())
        if message:
            details.append(message)

    for error_item in list(getattr(result, "errors", []) or []):
        message = _normalize_docling_message(getattr(error_item, "error_message", None) or str(error_item))
        module_name = _normalize_docling_message(getattr(error_item, "module_name", None))
        if module_name and message and not message.startswith(module_name):
            details.append(f"{module_name}: {message}")
        elif message:
            details.append(message)

    status = _normalize_docling_status(getattr(result, "status", None))
    return {
        "status": status,
        "details": _dedupe_messages(details),
    }


def _build_docling_failure_message(report: Dict[str, Any], extra_detail: Optional[str] = None) -> str:
    details = list(report.get("details", []) or [])
    if extra_detail:
        details = _dedupe_messages([extra_detail] + details)
    status = report.get("status")
    stage = _infer_docling_issue_stage(details)

    parts = ["Docling could not process this PDF into usable content."]
    if status and status != "success":
        parts.append(f"Docling status={status}.")
    if details:
        if stage != "docling":
            parts.append(f"Last reported stage={stage}.")
        parts.append(f"Detail: {details[0]}")
    parts.append(
        "Try running again, or retry with a lighter OCR configuration such as "
        "profile='balanced' and do_picture_description=False."
    )
    return " ".join(parts)


def _build_docling_recovery_warnings(
    report: Dict[str, Any],
    *,
    recovered_output: bool,
) -> List[Dict[str, Any]]:
    status = report.get("status")
    details = list(report.get("details", []) or [])
    has_issue = bool(details) or status in {"partial_success", "failure"}
    if not has_issue or not recovered_output:
        return []

    stage = _infer_docling_issue_stage(details)
    stage_phrase = f"{stage} " if stage != "docling" else ""
    warning: Dict[str, Any] = {
        "stage": stage,
        "severity": "warning",
        "code": "recoverable_docling_processing_issue",
        "message": (
            f"Docling reported an internal {stage_phrase}issue, but Navexa recovered content and continued. "
            "Extraction quality may be reduced."
        ),
    }
    if status:
        warning["docling_status"] = status
    if details:
        warning["details"] = details[:5]
    return [warning]


def _get_converter(
    *,
    do_ocr: bool,
    backend: str,
    force_full_page_ocr: bool,
    do_table_structure: bool,
    do_picture_description: bool,
    enable_remote_services: bool,
    quiet: bool,
) -> Any:
    from docling.datamodel.base_models import InputFormat  # type: ignore
    from docling.datamodel.pipeline_options import PdfPipelineOptions, RapidOcrOptions  # type: ignore
    from docling.document_converter import DocumentConverter, PdfFormatOption  # type: ignore

    backend_use = backend
    if backend_use == "onnxruntime" and importlib.util.find_spec("onnxruntime") is None:
        logging.getLogger("navexa").warning(
            "[docling] backend=onnxruntime requested but onnxruntime is not installed; falling back to backend=torch"
        )
        backend_use = "torch"

    key = (
        do_ocr,
        backend_use,
        force_full_page_ocr,
        do_table_structure,
        do_picture_description,
        enable_remote_services,
        quiet,
    )
    if key in _CONVERTER_CACHE:
        return _CONVERTER_CACHE[key]

    opts = PdfPipelineOptions()
    opts.do_ocr = do_ocr
    if do_ocr:
        ocr_options = RapidOcrOptions(backend=backend_use)  # type: ignore[arg-type]
        if hasattr(ocr_options, "force_full_page_ocr"):
            ocr_options.force_full_page_ocr = force_full_page_ocr
        if hasattr(ocr_options, "print_verbose"):
            ocr_options.print_verbose = not quiet
        if hasattr(ocr_options, "rapidocr_params") and quiet:
            rapidocr_params = dict(getattr(ocr_options, "rapidocr_params", {}) or {})
            rapidocr_params["Global.log_level"] = "WARNING"
            ocr_options.rapidocr_params = rapidocr_params
        opts.ocr_options = ocr_options
    if hasattr(opts, "do_table_structure"):
        opts.do_table_structure = do_table_structure
    if hasattr(opts, "do_picture_description"):
        opts.do_picture_description = do_picture_description
    if hasattr(opts, "enable_remote_services"):
        opts.enable_remote_services = enable_remote_services

    converter = DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=opts)})
    _CONVERTER_CACHE[key] = converter
    return converter


def _resolve_parser_model(parser_model: Optional[str]) -> str:
    resolved = (parser_model or os.getenv("NAVEXA_PARSER_MODEL", "docling")).strip().lower() or "docling"
    if resolved != "docling":
        raise ValueError(f"Unsupported parser_model '{resolved}'. Currently supported: docling")
    return resolved


def _resolve_output_format(output_format: Optional[str]) -> str:
    resolved = (output_format or os.getenv("NAVEXA_DOCLING_OUTPUT_FORMAT", "markdown")).strip().lower() or "markdown"
    if resolved not in {"markdown", "text"}:
        raise ValueError("output_format must be 'markdown' or 'text'")
    return resolved


def extract_docling_content(
    pdf_path: str,
    output_format: Optional[str] = None,
    parser_model: Optional[str] = None,
    docling_options: Optional[Dict[str, Any]] = None,
) -> Tuple[str, List[LayoutBlock], Dict[int, str], int, List[Dict[str, Any]]]:
    _safe_env()

    from docling_core.types.doc import ImageRefMode  # type: ignore

    _resolve_parser_model(parser_model)
    resolved_output_format = _resolve_output_format(output_format)
    options = _resolve_docling_parse_options(docling_options)
    _configure_docling_logging(quiet=options.quiet)

    logger = logging.getLogger("navexa")

    def _convert_and_collect(
        *,
        run_ocr: bool,
        run_force_full_page_ocr: bool,
    ) -> Tuple[Any, List[LayoutBlock], Dict[int, str], int, Dict[str, Any]]:
        converter = _get_converter(
            do_ocr=run_ocr,
            backend=options.backend,
            force_full_page_ocr=run_force_full_page_ocr,
            do_table_structure=options.do_table_structure,
            do_picture_description=options.do_picture_description,
            enable_remote_services=options.enable_remote_services,
            quiet=options.quiet,
        )
        _configure_docling_logging(quiet=options.quiet)
        with _capture_docling_pipeline_logs() as captured_records:
            try:
                result = converter.convert(pdf_path)
            except Exception as exc:
                detail = _normalize_docling_message(str(exc))
                report = {
                    "status": "failure",
                    "details": _dedupe_messages(
                        [_normalize_docling_message(record.getMessage()) for record in captured_records] + [detail]
                    ),
                }
                raise RuntimeError(_build_docling_failure_message(report)) from exc
        document_obj = getattr(result, "document", result)
        report = _build_docling_report(result, captured_records)

        extracted_blocks: List[LayoutBlock] = []
        page_text_map_local: Dict[int, List[str]] = {}
        page_order_local: Dict[int, int] = {}

        for item in list(getattr(document_obj, "texts", []) or []):
            text = (getattr(item, "text", None) or getattr(item, "orig", None) or "").strip()
            if not text:
                continue

            prov = list(getattr(item, "prov", []) or [])
            page_no = int(getattr(prov[0], "page_no", 1) or 1) if prov else 1
            order = page_order_local.get(page_no, 0)
            page_order_local[page_no] = order + 1

            label = getattr(item, "label", "text")
            block_type = getattr(label, "value", None) or str(label)
            label_name = getattr(label, "name", None) or (str(label).split(".")[-1] if str(label) else "TEXT")
            source_ref = getattr(item, "self_ref", None)
            provenance = _build_provenance_summary(prov, fallback_page=page_no)

            extracted_blocks.append(
                LayoutBlock(
                    block_id=f"d{page_no}_{order}",
                    page=page_no,
                    order=order,
                    text=text,
                    block_type=block_type,
                    label_name=str(label_name),
                    source_ref=str(source_ref) if source_ref is not None else None,
                    provenance=provenance,
                )
            )
            page_text_map_local.setdefault(page_no, []).append(text)

        extracted_blocks.sort(key=lambda b: (b.page, b.order))
        page_texts_local = {p: "\n".join(parts).strip() for p, parts in page_text_map_local.items()}
        page_count_local = len(getattr(document_obj, "pages", {}) or {})
        return document_obj, extracted_blocks, page_texts_local, page_count_local, report

    document, blocks, page_texts, num_pages, docling_report = _convert_and_collect(
        run_ocr=options.do_ocr,
        run_force_full_page_ocr=options.force_full_page_ocr,
    )

    # Image/scanned PDFs may still yield zero text; auto-retry once with OCR + full-page OCR.
    if not blocks and (not options.do_ocr or not options.force_full_page_ocr):
        logger.warning(
            "[docling] no text blocks in initial pass; retrying with OCR enabled (force_full_page_ocr=1)"
        )
        document, blocks, page_texts, num_pages, docling_report = _convert_and_collect(
            run_ocr=True,
            run_force_full_page_ocr=True,
        )

    if document is None:
        raise RuntimeError(_build_docling_failure_message(docling_report, extra_detail="Docling returned no document."))

    markdown_kwargs: Dict[str, object] = {}
    if options.image_mode == "none":
        markdown_kwargs["image_mode"] = ImageRefMode.PLACEHOLDER
        markdown_kwargs["image_placeholder"] = ""
    else:
        mode_map = {
            "placeholder": ImageRefMode.PLACEHOLDER,
            "embedded": ImageRefMode.EMBEDDED,
            "referenced": ImageRefMode.REFERENCED,
        }
        markdown_kwargs["image_mode"] = mode_map.get(options.image_mode, ImageRefMode.PLACEHOLDER)

    try:
        markdown = document.export_to_markdown(**markdown_kwargs)
    except Exception as exc:
        raise RuntimeError(
            _build_docling_failure_message(docling_report, extra_detail=f"Markdown export failed: {exc}")
        ) from exc
    if not markdown or not markdown.strip():
        raise RuntimeError(_build_docling_failure_message(docling_report, extra_detail="Docling returned empty markdown."))

    recovered_output = bool(markdown.strip()) or bool(blocks) or any(text.strip() for text in page_texts.values())
    docling_warnings = _build_docling_recovery_warnings(docling_report, recovered_output=recovered_output)

    # Never fail hard on empty blocks: synthesize minimal text blocks from extracted
    # page text or markdown so downstream indexing can continue.
    if not blocks:
        logger.warning("[docling] no text blocks after OCR; synthesizing fallback blocks from extracted markdown/text.")
        if not page_texts:
            fallback_text = _fallback_text_from_markdown(markdown) or "[NO_EXTRACTED_TEXT]"
            page_texts = {1: fallback_text}
            num_pages = max(1, int(num_pages or 0))
        blocks = _synthesize_blocks_from_text_map(page_texts)
        docling_warnings.append(
            {
                "stage": "docling",
                "severity": "warning",
                "code": "synthetic_docling_blocks",
                "message": (
                    "Docling did not return text blocks, so Navexa synthesized fallback blocks from extracted markdown/text. "
                    "Extraction quality may be reduced."
                ),
                "docling_status": docling_report.get("status"),
            }
        )

    if resolved_output_format == "text":
        output_text = "\n\n".join(page_texts[p] for p in sorted(page_texts) if page_texts[p].strip())
        if not output_text.strip():
            output_text = _fallback_text_from_markdown(markdown) or "[NO_EXTRACTED_TEXT]"
        return output_text, blocks, page_texts, num_pages, docling_warnings

    return markdown, blocks, page_texts, num_pages, docling_warnings
