from __future__ import annotations

import base64
import functools
import io
import json
import os
import tempfile
import threading
import time
import traceback
import uuid
from pathlib import Path
from typing import Any, Dict, Optional

import torch
from flask import Flask, jsonify, request, Response
import PIL.Image

from marker.config.parser import ConfigParser
from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
from marker.providers.registry import load_extensions
from marker.settings import settings as marker_settings

import request_store
import deepseek_ocr as dsocr

# ── Env configuration ────────────────────────────────────────────────
API_KEY = os.environ.get("API_KEY", "")
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11435")
DEESEEK_OCR_MODEL = os.environ.get("DEESEEK_OCR_MODEL", "deepseek-ocr")
OCR_BACKEND = os.environ.get("OCR_BACKEND", "marker")  # "marker" | "deepseek"
AMD_COMPUTE = os.environ.get("AMD_COMPUTE", "false").lower() in ("true", "1", "yes")
TORCH_DEVICE = os.environ.get("TORCH_DEVICE", "")
MODEL_DTYPE = os.environ.get("MODEL_DTYPE", "float32")

DEFAULT_LLM_SERVICE = os.environ.get("LLM_SERVICE", "marker.services.ollama.OllamaService")
DEFAULT_USE_LLM = os.environ.get("USE_LLM", "false").lower() in ("true", "1", "yes")
DEFAULT_OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "")
DEFAULT_OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
DEFAULT_OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "")

SUPPORTED_EXTENSIONS = set()
for provider_type in ("image", "pdf", "epub", "doc", "xls", "ppt"):
    SUPPORTED_EXTENSIONS.update(load_extensions(provider_type))
SUPPORTED_DISPLAY = sorted({ext.lstrip(".") for ext in SUPPORTED_EXTENSIONS})

_marker_dict: Optional[Dict[str, Any]] = None
_lock = threading.Lock()

BASE_URL = os.environ.get("SELF_URL", "http://localhost:8001")


def _configure_env():
    if AMD_COMPUTE and not TORCH_DEVICE:
        os.environ["TORCH_DEVICE"] = "cuda"
        os.environ["TORCH_DEVICE_MODEL"] = "cuda"
    if MODEL_DTYPE == "bfloat16":
        os.environ["MODEL_DTYPE"] = "bfloat16"
    if torch_device_override := os.environ.get("TORCH_DEVICE"):
        os.environ["TORCH_DEVICE"] = torch_device_override
        os.environ["TORCH_DEVICE_MODEL"] = torch_device_override


def get_model_dict() -> Dict[str, Any]:
    global _marker_dict
    if _marker_dict is None:
        with _lock:
            if _marker_dict is None:
                _marker_dict = create_model_dict()
    return _marker_dict


# ── Auth decorator ───────────────────────────────────────────────────
def require_api_key(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        if API_KEY:
            header_key = request.headers.get("X-API-Key", "")
            auth_header = request.headers.get("Authorization", "")
            if auth_header.startswith("Bearer "):
                header_key = header_key or auth_header[7:]
            if header_key != API_KEY:
                return jsonify({"success": False, "error": "Invalid or missing API key"}), 403
        return f(*args, **kwargs)
    return wrapper


# ── Option parsing ───────────────────────────────────────────────────
def parse_options(src: dict) -> dict:
    bool_fields = [
        "paginate_output", "force_ocr", "disable_image_extraction",
        "use_llm", "redo_inline_math", "strip_existing_ocr", "debug",
        "paginate", "skip_cache", "save_checkpoint",
        "disable_image_captions", "fence_synthetic_captions",
        "token_efficient_markdown", "add_block_ids",
        "include_markdown_in_chunks", "word_bboxes",
    ]
    str_fields = [
        "page_range", "output_format", "processors", "config_json",
        "converter_cls", "llm_service", "block_correction_prompt",
        "langs", "mode", "file_url", "additional_config",
        "extras", "pipeline_id", "webhook_url", "processing_location",
        "page_schema", "segmentation_schema", "model_override_settings",
    ]
    int_fields = ["max_pages", "workflowstepdata_id"]
    opts = {}
    for k in bool_fields:
        v = src.get(k)
        if isinstance(v, str):
            opts[k] = v.lower() in ("true", "1", "yes")
        elif v is not None:
            opts[k] = bool(v)
    for k in str_fields:
        v = src.get(k)
        if v is not None and v != "":
            opts[k] = str(v)
    for k in int_fields:
        v = src.get(k)
        if v is not None:
            opts[k] = int(v)
    return opts


def build_options(**extra: Any) -> Dict[str, Any]:
    opts: Dict[str, Any] = {}
    for k, v in extra.items():
        if v is not None:
            opts[k] = v
    opts.setdefault("output_format", "markdown")
    opts.setdefault("force_ocr", False)
    opts.setdefault("paginate_output", False)
    opts.setdefault("page_range", None)
    opts.setdefault("disable_multiprocessing", True)
    opts.setdefault("disable_image_extraction", False)
    opts.setdefault("output_dir", marker_settings.OUTPUT_DIR)
    opts.setdefault("llm_service", DEFAULT_LLM_SERVICE)
    if DEFAULT_OPENAI_BASE_URL:
        opts.setdefault("openai_base_url", DEFAULT_OPENAI_BASE_URL)
    if DEFAULT_OPENAI_API_KEY:
        opts.setdefault("openai_api_key", DEFAULT_OPENAI_API_KEY)
    if DEFAULT_OPENAI_MODEL:
        opts.setdefault("openai_model", DEFAULT_OPENAI_MODEL)
    if OLLAMA_HOST:
        opts.setdefault("ollama_base_url", OLLAMA_HOST)
    if DEESEEK_OCR_MODEL:
        opts.setdefault("ollama_model", DEESEEK_OCR_MODEL)
    config_parser = ConfigParser(opts)
    return config_parser.generate_config_dict()


# ── Conversion core ──────────────────────────────────────────────────
def convert_file_bytes(raw_bytes: bytes, filename: str, **opts: Any) -> dict:
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix="." + filename.rsplit(".", 1)[-1] if "." in filename else "")
    try:
        tmp.write(raw_bytes)
        tmp.close()
        filepath = tmp.name
        config_dict = build_options(**opts)
        config_dict["disable_tqdm"] = True
        model_dict = get_model_dict()
        parsed_opts = {k: v for k, v in opts.items() if v is not None}
        for key in ("output_format", "llm_service", "openai_base_url", "openai_api_key", "openai_model"):
            if key not in parsed_opts and key in config_dict:
                parsed_opts[key] = config_dict[key]
        parsed = ConfigParser(parsed_opts)
        converter = PdfConverter(
            config=config_dict,
            artifact_dict=model_dict,
            processor_list=parsed.get_processors(),
            renderer=parsed.get_renderer(),
            llm_service=parsed.get_llm_service(),
        )
        rendered = converter(filepath)
        text, _, images = text_from_rendered(rendered)
    except Exception as exc:
        traceback.print_exc()
        return {"success": False, "error": str(exc)}
    finally:
        tmp.close()
        if os.path.exists(tmp.name):
            os.remove(tmp.name)
    encoded_images: Dict[str, str] = {}
    img_fmt = marker_settings.OUTPUT_IMAGE_FORMAT
    for k, img in images.items():
        buf = io.BytesIO()
        img.save(buf, format=img_fmt)
        encoded_images[k] = base64.b64encode(buf.getvalue()).decode("utf-8")
    return {
        "format": opts.get("output_format", "markdown"),
        "output": text,
        "images_b64": encoded_images,
        "metadata": rendered.metadata if hasattr(rendered, 'metadata') else {},
        "success": True,
    }


def convert_with_ocr_backend(raw_bytes: bytes, filename: str, **opts: Any) -> dict:
    if OCR_BACKEND == "deepseek":
        page_range = opts.get("page_range")
        max_pages = opts.get("max_pages")
        ocr_result = dsocr.ocr_pdf(raw_bytes, max_pages=max_pages, page_range=page_range)
        return {
            "format": opts.get("output_format", "markdown"),
            "output": ocr_result.get("text", ""),
            "images_b64": {},
            "metadata": {"page_count": ocr_result.get("page_count", 0), "ocr_backend": "deepseek"},
            "success": ocr_result.get("success", False),
            "page_count": ocr_result.get("page_count", 0),
        }
    return convert_file_bytes(raw_bytes, filename, **opts)


# ── Background worker ────────────────────────────────────────────────
# The job queue (submit_job, _job_queue, _job_cond) lives in request_store.py
# so both the Flask app and gunicorn's post_fork hook can access it.


# ── Extraction / Segmentation helpers ────────────────────────────────
def _apply_extraction_schema(markdown_text: str, schema: dict) -> dict:
    result = {}
    for field_name, field_def in schema.items():
        if isinstance(field_def, dict):
            desc = field_def.get("description", "")
            ftype = field_def.get("type", "string")
            result[field_name] = _extract_field(markdown_text, field_name, desc, ftype)
        else:
            result[field_name] = _extract_field(markdown_text, field_name, str(field_def), "string")
    return result


def _extract_field(text: str, name: str, desc: str, ftype: str) -> Any:
    import re
    patterns = [
        re.compile(rf"(?:{re.escape(name)}\s*:?\s*(.*?))(?:\n\n|\Z)", re.IGNORECASE | re.DOTALL),
        re.compile(rf"(?:{re.escape(desc)}\s*:?\s*(.*?))(?:\n\n|\Z)", re.IGNORECASE | re.DOTALL),
    ]
    for pat in patterns:
        m = pat.search(text)
        if m:
            val = m.group(1).strip()
            if ftype == "number":
                try:
                    return float(val)
                except ValueError:
                    return val
            return val
    return None


def _apply_segmentation_schema(markdown_text: str, schema: dict) -> list[dict]:
    segments = []
    for seg_name, seg_desc in schema.items():
        if isinstance(seg_desc, str):
            desc_text = seg_desc
        elif isinstance(seg_desc, dict):
            desc_text = seg_desc.get("description", seg_name)
        else:
            desc_text = str(seg_desc)
        import re
        pattern = re.compile(rf"(?:{re.escape(desc_text)})(.*?)(?=(?:{'|'.join(re.escape(str(v)) if isinstance(v, str) else re.escape(str(v.get('description', ''))) for v in schema.values())})|\Z)", re.IGNORECASE | re.DOTALL)
        m = pattern.search(markdown_text)
        if m:
            segments.append({"name": seg_name, "description": desc_text, "content": m.group(1).strip()})
    return segments


def _extract_tables(conv_result: dict) -> list[dict]:
    output = conv_result.get("output", "")
    if conv_result.get("format") == "json":
        try:
            data = json.loads(output) if isinstance(output, str) else output
            if isinstance(data, dict):
                return data.get("tables", [])
            if isinstance(data, list):
                return [item for item in data if "table" in str(item).lower() or "rows" in str(item)]
        except (json.JSONDecodeError, TypeError):
            pass
    import re
    tables = []
    table_pattern = re.compile(r"\|(.+)\|[\s\S]*?(?=\n\n|\Z)", re.MULTILINE)
    for i, m in enumerate(table_pattern.finditer(output)):
        tables.append({"index": i, "table": m.group(0).strip()})
    return tables


# ── File management ──────────────────────────────────────────────────
def _file_storage_path() -> Path:
    p = Path(os.environ.get("UPLOAD_DIR", "/app/uploads"))
    p.mkdir(parents=True, exist_ok=True)
    return p


def _get_file_metadata(file_id: str) -> dict | None:
    meta_path = _file_storage_path() / file_id / "metadata.json"
    if not meta_path.exists():
        return None
    return json.loads(meta_path.read_text())


def _submit_and_poll(endpoint: str, raw_bytes: bytes, filename: str, form: dict | None = None, opts: dict | None = None) -> Response:
    if opts is None:
        opts = parse_options(form or {})
    request_id = request_store.submit_job(endpoint, raw_bytes, filename, opts)
    check_url = f"{BASE_URL}/api/v1/{endpoint}/{request_id}"
    return jsonify({
        "success": True,
        "request_id": request_id,
        "request_check_url": check_url,
    })


def _poll_result(endpoint: str, request_id: str) -> Response:
    entry = request_store.get_request(request_id)
    if entry is None:
        return jsonify({"success": False, "error": "Request not found"}), 404
    if entry["status"] == "processing":
        return jsonify({"status": "processing"})
    result = request_store.get_result(request_id)
    if result is None:
        return jsonify({"status": "complete", "success": False, "error": "Result not found"}), 500
    fmt = result.get("format", "markdown")
    payload = {
        "status": "complete",
        "success": result.get("success", False),
        "error": result.get("error"),
        "output_format": fmt,
    }
    output_text = result.get("output", "")
    if fmt in ("markdown", "html"):
        payload[fmt] = output_text
    elif fmt == "json":
        try:
            payload["json"] = json.loads(output_text) if output_text else None
        except (json.JSONDecodeError, TypeError):
            payload["json"] = output_text
    elif fmt == "chunks":
        payload["chunks"] = output_text
    payload["images"] = result.get("images_b64", {})
    payload["metadata"] = result.get("metadata", {})
    payload["page_count"] = result.get("page_count")
    return jsonify(payload)


# ── Flask app factory ────────────────────────────────────────────────
def create_app() -> Flask:
    app = Flask(__name__)

    # Start cleanup thread for expired results
    request_store.start_cleanup_thread()

    HTML_DOCS = r"""<html><head><title>marker-api</title>
<style>body{font-family:sans-serif;max-width:960px;margin:40px auto;padding:0 20px}
pre{background:#f4f4f4;padding:12px;border-radius:4px;overflow-x:auto}
table{border-collapse:collapse;width:100%;margin:16px 0}th,td{border:1px solid #ccc;padding:8px;text-align:left}
th{background:#f0f0f0}
.alert{padding:12px 16px;border-radius:4px;margin:16px 0}
.alert-deprecated{background:#fff3cd;border:1px solid #ffc107}
.alert-info{background:#d1ecf1;border:1px solid #bee5eb}
code{background:#f4f4f4;padding:2px 6px;border-radius:3px}
</style></head><body>
<h1>marker-api</h1>
<p>Datalab-compatible document conversion API. Convert PDFs, images, documents to Markdown/HTML/JSON.</p>

<div class="alert alert-info"><strong>OCR Backend:</strong> <code>{ocr_backend}</code></div>

<h2>Datalab-compatible Endpoints</h2>
<table>
<tr><th>Method</th><th>Path</th><th>Description</th></tr>
<tr><td>POST</td><td>/api/v1/convert</td><td>Convert document (async, returns request_id + check_url)</td></tr>
<tr><td>GET</td><td>/api/v1/convert/&lt;id&gt;</td><td>Poll conversion result</td></tr>
<tr><td>POST</td><td>/api/v1/extract</td><td>Structured extraction with JSON schema</td></tr>
<tr><td>GET</td><td>/api/v1/extract/&lt;id&gt;</td><td>Poll extraction result</td></tr>
<tr><td>POST</td><td>/api/v1/segment</td><td>Document segmentation with JSON schema</td></tr>
<tr><td>GET</td><td>/api/v1/segment/&lt;id&gt;</td><td>Poll segmentation result</td></tr>
<tr><td>POST</td><td>/api/v1/ocr</td><td>[DEPRECATED] OCR-only extraction</td></tr>
<tr><td>GET</td><td>/api/v1/ocr/&lt;id&gt;</td><td>Poll OCR result</td></tr>
<tr><td>POST</td><td>/api/v1/table_rec</td><td>[DEPRECATED] Table recognition</td></tr>
<tr><td>GET</td><td>/api/v1/table_rec/&lt;id&gt;</td><td>Poll table recognition result</td></tr>
<tr><td>POST</td><td>/api/v1/marker</td><td>[DEPRECATED] Sync conversion (returns directly)</td></tr>
<tr><td>POST</td><td>/api/v1/create-document</td><td>Create DOCX from markdown</td></tr>
</table>

<h3>File Management</h3>
<table>
<tr><th>Method</th><th>Path</th><th>Description</th></tr>
<tr><td>POST</td><td>/api/v1/files/upload</td><td>Request file upload URL</td></tr>
<tr><td>GET</td><td>/api/v1/files</td><td>List uploaded files</td></tr>
<tr><td>GET</td><td>/api/v1/files/&lt;id&gt;</td><td>Get file metadata</td></tr>
<tr><td>GET</td><td>/api/v1/files/&lt;id&gt;/download</td><td>Get file download URL</td></tr>
<tr><td>DELETE</td><td>/api/v1/files/&lt;id&gt;</td><td>Delete file</td></tr>
</table>

<h3>Legacy Endpoints</h3>
<table>
<tr><th>Method</th><th>Path</th><th>Description</th></tr>
<tr><td>POST</td><td>/marker</td><td>Sync conversion (original)</td></tr>
<tr><td>POST</td><td>/v1/conversions</td><td>Sync conversion (original)</td></tr>
<tr><td>POST</td><td>/v1/files/convert</td><td>Sync conversion via base64 JSON</td></tr>
</table>

<h2>Authentication</h2>
<p>All <code>/api/v1/*</code> endpoints accept <code>X-API-Key</code> header (or <code>Authorization: Bearer</code>).</p>

<h2>POST /api/v1/convert</h2>
<pre><code>curl -X POST http://localhost:8001/api/v1/convert \\
  -H "X-API-Key: YOUR_KEY" \\
  -F "file=@document.pdf" \\
  -F "output_format=markdown" \\
  -F "mode=balanced"</code></pre>

<h2>Poll for results</h2>
<pre><code>curl http://localhost:8001/api/v1/convert/REQUEST_ID \\
  -H "X-API-Key: YOUR_KEY"</code></pre>

<h3>Parameters</h3>
<table>
<tr><th>Parameter</th><th>Type</th><th>Default</th><th>Description</th></tr>
<tr><td><b>file</b></td><td>file</td><td>required</td><td>Document to convert</td></tr>
<tr><td>output_format</td><td>string</td><td>markdown</td><td>markdown, html, json, chunks</td></tr>
<tr><td>mode</td><td>string</td><td>fast</td><td>fast, balanced, accurate</td></tr>
<tr><td>page_range</td><td>string</td><td>all</td><td>e.g. "0,5-10"</td></tr>
<tr><td>max_pages</td><td>int</td><td>all</td><td>Max pages to process</td></tr>
<tr><td>paginate</td><td>bool</td><td>false</td><td>Add page delimiters</td></tr>
</table>

<h3>Environment</h3>
<table>
<tr><th>Variable</th><th>Default</th><th>Description</th></tr>
<tr><td>OCR_BACKEND</td><td>marker</td><td>"marker" (marker-pdf) or "deepseek" (deepseek-ocr via ollama)</td></tr>
<tr><td>API_KEY</td><td>(none)</td><td>API key for X-API-Key auth (empty = disabled)</td></tr>
<tr><td>DEEPSEEK_OLLAMA_HOST</td><td>http://localhost:11434</td><td>Ollama host for deepseek-ocr</td></tr>
<tr><td>DEEPSEEK_OCR_MODEL</td><td>deepseek-ocr:latest</td><td>Model name for deepseek-ocr</td></tr>
<tr><td>OLLAMA_HOST</td><td>http://localhost:11435</td><td>Ollama host for LLM correction</td></tr>
<tr><td>DEESEEK_OCR_MODEL</td><td>deepseek-ocr</td><td>LLM correction model name</td></tr>
<tr><td>TORCH_DEVICE</td><td>auto</td><td>cpu, cuda, rocm</td></tr>
<tr><td>PORT</td><td>8001</td><td>Listening port</td></tr>
<tr><td>GUNICORN_WORKERS</td><td>3</td><td>Number of workers</td></tr>
</table>
<h3>Supported Formats</h3>
<p><code>{formats}</code></p>
</body></html>"""

    # ── Legacy /info endpoints ───────────────────────────────────────
    @app.route("/")
    def docs():
        body = HTML_DOCS.replace("{ocr_backend}", OCR_BACKEND).replace("{formats}", ", ".join(SUPPORTED_DISPLAY))
        return body

    @app.route("/health")
    def health():
        try:
            device = TORCH_DEVICE or (marker_settings.TORCH_DEVICE_MODEL if hasattr(marker_settings, 'TORCH_DEVICE_MODEL') else 'auto')
        except Exception:
            device = "unknown"
        return jsonify({
            "status": "ok",
            "ollama": OLLAMA_HOST,
            "ocr_backend": OCR_BACKEND,
            "amd_compute": AMD_COMPUTE,
            "torch_device": device,
            "supported_formats": SUPPORTED_DISPLAY,
            "provider": "flask",
            "default_llm_service": DEFAULT_LLM_SERVICE,
            "default_use_llm": DEFAULT_USE_LLM,
            "openai_base_url": DEFAULT_OPENAI_BASE_URL or None,
            "openai_model": DEFAULT_OPENAI_MODEL or None,
        })

    # ── Existing original endpoints (unchanged behavior) ────────────
    @app.route("/marker", methods=["POST"])
    def convert_sync():
        if "file" in request.files:
            file = request.files["file"]
            filename = file.filename or "file"
            raw = file.read()
            fmt = request.form.get("output_format", "markdown")
            opts = parse_options(request.form)
            opts["output_format"] = fmt
            result = convert_with_ocr_backend(raw, filename, **opts)
            if not result["success"]:
                return jsonify(result), 500
            if fmt == "markdown":
                return Response(
                    result["output"],
                    mimetype="text/plain",
                    headers={"Content-Disposition": f'attachment; filename="{filename.rsplit(".", 1)[0]}.md"'},
                )
            return jsonify(result)
        if request.is_json:
            data = request.get_json()
            if "file_b64" not in data or not data.get("filename"):
                return jsonify({"error": "JSON body must include 'file_b64' and 'filename'"}), 400
            raw = base64.b64decode(data["file_b64"])
            fmt = data.get("output_format", "markdown")
            opts = parse_options(data)
            opts["output_format"] = fmt
            result = convert_with_ocr_backend(raw, data["filename"], **opts)
            return jsonify(result)
        return jsonify({"error": "No file provided. Use multipart/form-data or JSON with 'file_b64'."}), 400

    @app.route("/v1/conversions", methods=["POST"])
    def convert_async_style():
        if request.is_json:
            data = request.get_json()
            if "file_b64" not in data or not data.get("filename"):
                return jsonify({"error": "JSON body must include 'file_b64' and 'filename'"}), 400
            raw = base64.b64decode(data["file_b64"])
            fmt = data.get("output_format", "markdown")
            opts = parse_options(data)
            opts["output_format"] = fmt
            result = convert_with_ocr_backend(raw, data["filename"], **opts)
            if not result["success"]:
                return jsonify(result), 500
            return jsonify({
                "id": str(uuid.uuid4()),
                "filename": data["filename"],
                "format": fmt,
                "output": result["output"],
                "success": True,
                "images_b64": result.get("images_b64", {}),
                "metadata": result.get("metadata", {}),
            })
        if "file" in request.files:
            file = request.files["file"]
            filename = file.filename or "file"
            raw = file.read()
            fmt = request.form.get("output_format", "markdown")
            opts = parse_options(request.form)
            opts["output_format"] = fmt
            result = convert_with_ocr_backend(raw, filename, **opts)
            if not result["success"]:
                return jsonify(result), 500
            return jsonify({
                "id": str(uuid.uuid4()),
                "filename": filename,
                "format": fmt,
                "output": result["output"],
                "success": True,
                "images_b64": result.get("images_b64", {}),
                "metadata": result.get("metadata", {}),
            })
        return jsonify({"error": "No file provided."}), 400

    @app.route("/v1/files/convert", methods=["POST"])
    def convert_files():
        if "file" in request.files:
            file = request.files["file"]
            filename = file.filename or "file"
            raw = file.read()
            fmt = request.form.get("output_format", "markdown")
            opts = parse_options(request.form)
            opts["output_format"] = fmt
            result = convert_with_ocr_backend(raw, filename, **opts)
            filename_var = filename
        elif request.is_json:
            data = request.get_json()
            if "file_b64" not in data or not data.get("filename"):
                return jsonify({"error": "JSON body must include 'file_b64' and 'filename'"}), 400
            raw = base64.b64decode(data["file_b64"])
            fmt = data.get("output_format", "markdown")
            opts = parse_options(data)
            opts["output_format"] = fmt
            result = convert_with_ocr_backend(raw, data["filename"], **opts)
            filename_var = data["filename"]
        else:
            return jsonify({"error": "No file provided."}), 400
        if not result["success"]:
            return jsonify(result), 500
        return jsonify({
            "filename": filename_var,
            "format": fmt,
            "output": result["output"],
            "images_b64": result.get("images_b64", {}),
            "metadata": result.get("metadata", {}),
        })

    # ── New Datalab-compatible endpoints ─────────────────────────────

    # POST /api/v1/convert — async submit
    @app.route("/api/v1/convert", methods=["POST"])
    @require_api_key
    def api_convert_submit():
        if "file" in request.files:
            file = request.files["file"]
            raw = file.read()
            filename = file.filename or "file"
        elif request.is_json:
            data = request.get_json()
            if "file_b64" in data:
                raw = base64.b64decode(data["file_b64"])
                filename = data.get("filename", "file")
            elif "file_url" in data:
                import requests as req
                resp = req.get(data["file_url"], timeout=120)
                resp.raise_for_status()
                raw = resp.content
                filename = data.get("filename", data["file_url"].rsplit("/", 1)[-1] or "file")
            else:
                return jsonify({"error": "Provide file, file_b64, or file_url"}), 400
        else:
            return jsonify({"error": "No file provided"}), 400
        source = request.form if request.form else (request.get_json() if request.is_json else {})
        opts = parse_options(source)
        opts["output_format"] = source.get("output_format", "markdown")
        return _submit_and_poll("convert", raw, filename, opts=opts)

    # GET /api/v1/convert/<request_id> — poll
    @app.route("/api/v1/convert/<request_id>", methods=["GET"])
    @require_api_key
    def api_convert_poll(request_id: str):
        return _poll_result("convert", request_id)

    # POST /api/v1/marker — deprecated (Datalab compat: async submit-and-poll)
    @app.route("/api/v1/marker", methods=["POST"])
    @require_api_key
    def api_marker_deprecated():
        if "file" in request.files:
            file = request.files["file"]
            raw = file.read()
            filename = file.filename or "file"
        elif request.is_json:
            data = request.get_json()
            if "file_b64" in data:
                raw = base64.b64decode(data["file_b64"])
                filename = data.get("filename", "file")
            else:
                return jsonify({"error": "Provide file or file_b64"}), 400
        else:
            return jsonify({"error": "No file provided"}), 400
        source = request.form if request.form else (request.get_json() if request.is_json else {})
        opts = parse_options(source)
        opts.setdefault("output_format", "markdown")
        resp = _submit_and_poll("convert", raw, filename, opts=opts)
        if isinstance(resp, Response):
            resp.headers["Warning"] = "299 marker-api \"POST /api/v1/marker is deprecated, use /api/v1/convert\""
        return resp

    # POST /api/v1/extract — structured extraction
    @app.route("/api/v1/extract", methods=["POST"])
    @require_api_key
    def api_extract_submit():
        if "file" in request.files:
            file = request.files["file"]
            raw = file.read()
            filename = file.filename or "file"
        elif request.is_json:
            data = request.get_json()
            if "file_b64" in data:
                raw = base64.b64decode(data["file_b64"])
                filename = data.get("filename", "file")
            else:
                return jsonify({"error": "Provide file or file_b64"}), 400
        else:
            return jsonify({"error": "No file provided"}), 400
        source = request.form if request.form else (request.get_json() if request.is_json else {})
        opts = parse_options(source)
        opts.setdefault("output_format", "markdown")
        return _submit_and_poll("extract", raw, filename, opts=opts)

    # GET /api/v1/extract/<request_id>
    @app.route("/api/v1/extract/<request_id>", methods=["GET"])
    @require_api_key
    def api_extract_poll(request_id: str):
        return _poll_result("extract", request_id)

    # POST /api/v1/segment — document segmentation
    @app.route("/api/v1/segment", methods=["POST"])
    @require_api_key
    def api_segment_submit():
        if "file" in request.files:
            file = request.files["file"]
            raw = file.read()
            filename = file.filename or "file"
        elif request.is_json:
            data = request.get_json()
            if "file_b64" in data:
                raw = base64.b64decode(data["file_b64"])
                filename = data.get("filename", "file")
            else:
                return jsonify({"error": "Provide file or file_b64"}), 400
        else:
            return jsonify({"error": "No file provided"}), 400
        source = request.form if request.form else (request.get_json() if request.is_json else {})
        opts = parse_options(source)
        opts.setdefault("output_format", "markdown")
        return _submit_and_poll("segment", raw, filename, opts=opts)

    # GET /api/v1/segment/<request_id>
    @app.route("/api/v1/segment/<request_id>", methods=["GET"])
    @require_api_key
    def api_segment_poll(request_id: str):
        return _poll_result("segment", request_id)

    # POST /api/v1/ocr — deprecated OCR-only endpoint
    @app.route("/api/v1/ocr", methods=["POST"])
    @require_api_key
    def api_ocr_submit():
        if "file" not in request.files:
            return jsonify({"error": "No file provided"}), 400
        file = request.files["file"]
        raw = file.read()
        filename = file.filename or "file"
        opts = parse_options(request.form)
        return _submit_and_poll("ocr", raw, filename, opts=opts)

    # GET /api/v1/ocr/<request_id>
    @app.route("/api/v1/ocr/<request_id>", methods=["GET"])
    @require_api_key
    def api_ocr_poll(request_id: str):
        return _poll_result("ocr", request_id)

    # POST /api/v1/table_rec — deprecated table recognition
    @app.route("/api/v1/table_rec", methods=["POST"])
    @require_api_key
    def api_table_rec_submit():
        if "file" not in request.files:
            return jsonify({"error": "No file provided"}), 400
        file = request.files["file"]
        raw = file.read()
        filename = file.filename or "file"
        opts = parse_options(request.form)
        return _submit_and_poll("table_rec", raw, filename, opts=opts)

    # GET /api/v1/table_rec/<request_id>
    @app.route("/api/v1/table_rec/<request_id>", methods=["GET"])
    @require_api_key
    def api_table_rec_poll(request_id: str):
        return _poll_result("table_rec", request_id)

    # POST /api/v1/create-document — create DOCX from markdown
    @app.route("/api/v1/create-document", methods=["POST"])
    @require_api_key
    def api_create_document():
        data = request.get_json() if request.is_json else {}
        if not data or "markdown" not in data:
            return jsonify({"error": "JSON body must include 'markdown'"}), 400
        md = data["markdown"]
        output_format = data.get("output_format", "docx")
        try:
            import subprocess
            import base64
            with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as f:
                f.write(md)
                md_path = f.name
            out_path = md_path.replace(".md", f".{output_format}")
            if output_format == "docx":
                subprocess.run(["pandoc", md_path, "-o", out_path], capture_output=True, timeout=60)
                with open(out_path, "rb") as f:
                    b64 = base64.b64encode(f.read()).decode()
                os.unlink(out_path)
                os.unlink(md_path)
                return jsonify({
                    "success": True,
                    "output_format": "docx",
                    "output_base64": b64,
                })
            else:
                return jsonify({"error": f"Unsupported output format: {output_format}"}), 400
        except FileNotFoundError:
            return jsonify({"error": "pandoc not installed, cannot create document"}), 500
        except Exception as exc:
            return jsonify({"success": False, "error": str(exc)}), 500

    # ── File management endpoints ────────────────────────────────────
    @app.route("/api/v1/files/upload", methods=["POST"])
    @require_api_key
    def api_file_upload():
        data = request.get_json() if request.is_json else {}
        filename = data.get("filename", "file")
        content_type = data.get("content_type", "application/octet-stream")
        import time as tmod
        file_id = str(uuid.uuid4())
        file_dir = _file_storage_path() / file_id
        file_dir.mkdir(parents=True, exist_ok=True)
        metadata = {
            "file_id": file_id,
            "original_filename": filename,
            "filename": filename,
            "content_type": content_type,
            "reference": f"datalab://file-{file_id}",
            "created_at": time.time(),
            "created": tmod.strftime("%Y-%m-%dT%H:%M:%SZ", tmod.gmtime()),
            "file_size": None,
            "upload_status": "pending",
            "error": None,
        }
        (file_dir / "metadata.json").write_text(json.dumps(metadata))
        import secrets
        token = secrets.token_urlsafe(32)
        upload_url = f"{BASE_URL}/api/v1/files/{file_id}/upload/{token}"
        (file_dir / "upload_token.txt").write_text(token)
        return jsonify({
            "file_id": file_id,
            "upload_url": upload_url,
            "reference": f"datalab://file-{file_id}",
        })

    @app.route("/api/v1/files/<file_id>/upload/<token>", methods=["PUT"])
    def api_file_upload_put(file_id: str, token: str):
        file_dir = _file_storage_path() / file_id
        token_path = file_dir / "upload_token.txt"
        if not token_path.exists() or token_path.read_text().strip() != token:
            return jsonify({"error": "Invalid upload token"}), 403
        (file_dir / "content").write_bytes(request.data)
        metadata_path = file_dir / "metadata.json"
        if metadata_path.exists():
            import time as tmod
            meta = json.loads(metadata_path.read_text())
            meta["uploaded_at"] = tmod.time()
            meta["file_size"] = len(request.data)
            meta["upload_status"] = "completed"
            metadata_path.write_text(json.dumps(meta))
        return jsonify({"success": True})

    @app.route("/api/v1/files", methods=["GET"])
    @require_api_key
    def api_file_list():
        storage = _file_storage_path()
        limit = int(request.args.get("limit", 50))
        offset = int(request.args.get("offset", 0))
        files = []
        for child in sorted(storage.iterdir(), reverse=True):
            if child.is_dir():
                meta = _get_file_metadata(child.name)
                if meta:
                    files.append(meta)
        return jsonify({"files": files[offset:offset + limit], "total": len(files)})

    @app.route("/api/v1/files/<file_id>", methods=["GET"])
    @require_api_key
    def api_file_get(file_id: str):
        meta = _get_file_metadata(file_id)
        if meta is None:
            return jsonify({"error": "File not found"}), 404
        return jsonify(meta)

    @app.route("/api/v1/files/<file_id>/confirm", methods=["GET"])
    @require_api_key
    def api_file_confirm(file_id: str):
        meta = _get_file_metadata(file_id)
        if meta is None:
            return jsonify({"error": "File not found"}), 404
        file_dir = _file_storage_path() / file_id
        content_path = file_dir / "content"
        if not content_path.exists():
            return jsonify({"error": "File content not uploaded yet"}), 400
        return jsonify({"success": True, "file_id": file_id})

    @app.route("/api/v1/files/<file_id>/download", methods=["GET"])
    @require_api_key
    def api_file_download(file_id: str):
        meta = _get_file_metadata(file_id)
        if meta is None:
            return jsonify({"error": "File not found"}), 404
        expires_in = int(request.args.get("expires_in", 3600))
        import secrets
        token = secrets.token_urlsafe(32)
        file_dir = _file_storage_path() / file_id
        (file_dir / "download_token.txt").write_text(token)
        download_url = f"{BASE_URL}/api/v1/files/{file_id}/download/{token}?expires_in={expires_in}"
        return jsonify({"download_url": download_url})

    @app.route("/api/v1/files/<file_id>/download/<token>", methods=["GET"])
    def api_file_download_token(file_id: str, token: str):
        file_dir = _file_storage_path() / file_id
        token_path = file_dir / "download_token.txt"
        if not token_path.exists() or token_path.read_text().strip() != token:
            return jsonify({"error": "Invalid download token"}), 403
        content_path = file_dir / "content"
        if not content_path.exists():
            return jsonify({"error": "File content not found"}), 404
        meta = _get_file_metadata(file_id) or {}
        return Response(
            content_path.read_bytes(),
            mimetype=meta.get("content_type", "application/octet-stream"),
            headers={"Content-Disposition": f'attachment; filename="{meta.get("filename", "file")}"'},
        )

    @app.route("/api/v1/files/<file_id>", methods=["DELETE"])
    @require_api_key
    def api_file_delete(file_id: str):
        file_dir = _file_storage_path() / file_id
        if file_dir.exists():
            import shutil
            shutil.rmtree(str(file_dir))
        return jsonify({"success": True})

    # ── Thumbnails ───────────────────────────────────────────────────
    @app.route("/api/v1/thumbnails/<lookup_key>", methods=["GET"])
    @require_api_key
    def api_thumbnails(lookup_key: str):
        thumb_width = int(request.args.get("thumb_width", 300))
        page_range = request.args.get("page_range")
        entry = request_store.get_request(lookup_key)
        if entry is None:
            return jsonify({"success": False, "error": "Request not found"}), 404
        result = request_store.get_result(lookup_key)
        if result is None:
            return jsonify({"success": False, "error": "No result found"}), 404
        metadata = result.get("metadata", {})
        page_count = metadata.get("page_count", 0) if isinstance(metadata, dict) else 0
        if page_range:
            pages = []
            for part in page_range.split(","):
                part = part.strip()
                if "-" in part:
                    a, b = part.split("-", 1)
                    pages.extend(range(int(a.strip()), int(b.strip()) + 1))
                else:
                    pages.append(int(part))
        else:
            pages = list(range(page_count))
        thumbnails = []
        for p in pages:
            if p < page_count:
                thumbnails.append("")  # placeholder — no rendered page images stored
        return jsonify({
            "success": True,
            "thumbnails": thumbnails,
        })

    return app


app_instance = create_app()

if __name__ == "__main__":
    _configure_env()
    _marker_dict = create_model_dict()
    port = int(os.environ.get("PORT", "8001"))
    host = os.environ.get("HOST", "0.0.0.0")
    print("=" * 60)
    print("marker-api starting")
    print(f"  OCR_BACKEND       = {OCR_BACKEND}")
    print(f"  API_KEY           = {'set' if API_KEY else '(not set)'}")
    print(f"  OLLAMA_HOST       = {OLLAMA_HOST}")
    print(f"  TORCH_DEVICE      = {TORCH_DEVICE or '(auto)'}")
    print(f"  LISTENING ON      = {host}:{port}")
    print(f"  FORMATS           = {', '.join(SUPPORTED_DISPLAY)}")
    print("=" * 60)
    app_instance.run(host=host, port=port, debug=(os.environ.get("FLASK_DEBUG", "0") == "1"), threaded=True)