"""Phase 0 user-provided material ingestion."""

from __future__ import annotations

import base64
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import requests


DEFAULT_FIRERED_OCR_ENDPOINT = "http://192.168.50.100:8001"
DEFAULT_OCR_MAX_PAGES = 50


@dataclass(frozen=True)
class OcrResult:
    text: str
    pages_processed: int
    output_path: Path


def safe_filename(path: Path) -> str:
    name = path.name.strip()
    return name or "material"


def extract_pdf_text(path: Path) -> tuple[str, int, bool]:
    from pypdf import PdfReader

    reader = PdfReader(str(path))
    chunks: list[str] = []
    for index, page in enumerate(reader.pages, start=1):
        text = (page.extract_text() or "").strip()
        if text:
            chunks.append(f"\n\n## Page {index}\n\n{text}")
    combined = "".join(chunks).strip()
    ocr_required = len(combined) < max(20, len(reader.pages) * 20)
    return combined, len(reader.pages), ocr_required


def ocr_endpoint_from_env() -> str:
    return os.environ.get("DEEP_RESEARCH_OCR_ENDPOINT", DEFAULT_FIRERED_OCR_ENDPOINT).rstrip("/")


def ocr_max_pages_from_env() -> int:
    raw = os.environ.get("DEEP_RESEARCH_OCR_MAX_PAGES")
    if not raw:
        return DEFAULT_OCR_MAX_PAGES
    try:
        return max(1, int(raw))
    except ValueError:
        return DEFAULT_OCR_MAX_PAGES


def render_pdf_pages(pdf_path: Path, output_dir: Path, *, max_pages: int) -> list[Path]:
    import fitz

    pages_dir = output_dir / f"{pdf_path.stem}.ocr-pages"
    pages_dir.mkdir(parents=True, exist_ok=True)
    image_paths: list[Path] = []
    doc = fitz.open(pdf_path)
    try:
        for index, page in enumerate(doc[:max_pages], start=1):
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
            image_path = pages_dir / f"page-{index:03d}.png"
            pix.save(image_path)
            image_paths.append(image_path)
    finally:
        doc.close()
    return image_paths


def data_url_for_image(path: Path) -> str:
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def call_firered_ocr(image_path: Path, *, endpoint: str) -> str:
    payload = {
        "model": "firered-ocr",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "请识别图片中的全部文字，保持原有顺序，只输出文字。"},
                    {"type": "image_url", "image_url": {"url": data_url_for_image(image_path)}},
                ],
            }
        ],
        "temperature": 0,
        "max_tokens": 3000,
    }
    response = requests.post(f"{endpoint.rstrip('/')}/v1/chat/completions", json=payload, timeout=60)
    if not response.ok:
        raise RuntimeError(f"{response.status_code} {response.text[:500]}")
    data = response.json()
    return str(data["choices"][0]["message"].get("content") or "").strip()


def ocr_pdf_with_firered(*, pdf_path: Path, output_dir: Path, endpoint: str, max_pages: int) -> OcrResult:
    image_paths = render_pdf_pages(pdf_path, output_dir, max_pages=max_pages)
    chunks: list[str] = []
    for index, image_path in enumerate(image_paths, start=1):
        text = call_firered_ocr(image_path, endpoint=endpoint)
        if text:
            chunks.append(f"\n\n## OCR Page {index}\n\n{text}")
    combined = "".join(chunks).strip()
    output_path = output_dir / f"{pdf_path.stem}.ocr.md"
    body = [
        f"# OCR Material: {pdf_path.name}",
        "",
        f"- source_path: {pdf_path}",
        f"- endpoint: {endpoint}",
        f"- pages_processed: {len(image_paths)}",
        "",
        combined or "OCR 未返回可用文本。",
        "",
    ]
    output_path.write_text("\n".join(body), encoding="utf-8")
    return OcrResult(text=combined, pages_processed=len(image_paths), output_path=output_path)


def ingest_input_materials(project_root: Path, materials: list[str] | None) -> list[dict[str, Any]]:
    inventory: list[dict[str, Any]] = []
    if not materials:
        return inventory

    inputs_dir = project_root / "phase0" / "inputs"
    extracted_dir = project_root / "phase0" / "extracted"
    ocr_endpoint = ocr_endpoint_from_env()
    ocr_max_pages = ocr_max_pages_from_env()
    inputs_dir.mkdir(parents=True, exist_ok=True)
    extracted_dir.mkdir(parents=True, exist_ok=True)

    for raw in materials:
        source = Path(raw).expanduser()
        if not source.exists():
            inventory.append({"kind": "note", "note": raw})
            continue

        copied = inputs_dir / safe_filename(source)
        shutil.copy2(source, copied)
        item: dict[str, Any] = {
            "kind": source.suffix.lower().lstrip(".") or "file",
            "source_path": str(source),
            "copied_to": str(copied.relative_to(project_root)),
            "size_bytes": source.stat().st_size,
        }

        if source.suffix.lower() == ".pdf":
            text, pages, ocr_required = extract_pdf_text(source)
            extracted = extracted_dir / f"{source.stem}.md"
            ocr_result: OcrResult | None = None
            ocr_error: str | None = None
            if ocr_required:
                try:
                    ocr_result = ocr_pdf_with_firered(
                        pdf_path=source,
                        output_dir=extracted_dir,
                        endpoint=ocr_endpoint,
                        max_pages=min(pages, ocr_max_pages),
                    )
                    if ocr_result.text:
                        text = "\n\n".join(part for part in [text, ocr_result.text] if part)
                except Exception as exc:  # noqa: BLE001 - ingestion should not block project init.
                    ocr_error = str(exc)

            body = [
                f"# Extracted Material: {source.name}",
                "",
                f"- source_path: {source}",
                f"- copied_to: {copied.relative_to(project_root)}",
                f"- pages: {pages}",
                f"- ocr_required: {str(ocr_required).lower()}",
                f"- ocr_status: {'completed' if ocr_result else 'failed' if ocr_error else 'not_required'}",
                "",
                text or "未能从 PDF 直接抽取文本；该材料可能需要 OCR。",
                "",
            ]
            if ocr_error:
                body.extend(["## OCR Error", "", ocr_error, ""])
            extracted.write_text("\n".join(body), encoding="utf-8")
            item.update(
                {
                    "pages": pages,
                    "extracted_to": str(extracted.relative_to(project_root)),
                    "text_chars": len(text),
                    "ocr_required": ocr_required,
                    "ocr_status": "completed" if ocr_result else "failed" if ocr_error else "not_required",
                }
            )
            if ocr_result:
                item.update(
                    {
                        "ocr_endpoint": ocr_endpoint,
                        "ocr_pages_processed": ocr_result.pages_processed,
                        "ocr_extracted_to": str(ocr_result.output_path.relative_to(project_root)),
                        "ocr_text_chars": len(ocr_result.text),
                    }
                )
            if ocr_error:
                item["ocr_error"] = ocr_error
        else:
            item["ocr_required"] = source.suffix.lower() in {".png", ".jpg", ".jpeg", ".tif", ".tiff"}
        inventory.append(item)

    return inventory


def render_material_inventory(inventory: list[dict[str, Any]]) -> str:
    if not inventory:
        return "- 暂无；可通过 `--input-material` 加入审计报告、问题清单或内部记录。"
    lines: list[str] = []
    for item in inventory:
        if item.get("kind") == "note":
            lines.append(f"- 备注：{item.get('note', '')}")
            continue
        marker = "；需要 OCR" if item.get("ocr_required") else ""
        ocr = f"；OCR：{item.get('ocr_status')}" if item.get("ocr_status") else ""
        extracted = item.get("extracted_to")
        extra = f"；抽取文本：{extracted}" if extracted else ""
        lines.append(
            f"- {item.get('copied_to')}（{item.get('kind')}，{item.get('size_bytes', 0)} bytes{extra}{marker}{ocr}）"
        )
    return "\n".join(lines)
