from __future__ import annotations

import base64
import io
import json
import os
import tempfile
import traceback
import uuid
from typing import Any, Dict, Optional

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

# Marker imports
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

# Ollama / OCR fallback configuration
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://10.0.1.127:11434")
DEESEEK_OCR_MODEL = os.environ.get("DEESEEK_OCR_MODEL", "deepseek-ocr")

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")

_marker_dict: Optional[Dict[str, Any]] = None

# Collect supported file extensions
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})

# Env-var overrides for default LLM service
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", "")


def _configure_env():
    """Apply AMD GPU / environment overrides before torch loads."""
    if AMD_COMPUTE and not TORCH_DEVICE:
        os.environ["TORCH_DEVICE"] = "cuda"
        os.environ["TORCH_DEVICE_MODEL"] = "cuda"
        os.environ["HSA_OVERRIDE_GFX_VERSION"] = "9.0.6"
    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:
        _marker_dict = create_model_dict()
    return _marker_dict


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

    # Ensure defaults are set
    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)

    # Pass default LLM settings from environment if not explicitly provided
    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()


def convert_file_bytes(raw_bytes: bytes, filename: str, **opts: Any) -> dict:
    """Convert in-memory file bytes to markdown (or other format)."""

    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()

        # Build parsed options, preserving server default LLM config when client omits it
        parsed_opts = {k: v for k, v in opts.items() if v is not None}
        for key in ("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 'rendered' in dir() else {},
        "success": True,
    }


def create_app() -> Flask:
    app = Flask(__name__)

    # ---- docs page ----
    HTML_DOCS = r"""<html><head><title>marker-api</title>
<style>body{font-family:sans-serif;max-width:900px;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}</style></head><body>
<h1>marker-api</h1>
<p>Convert PDFs, EPUBs, DOCX, XLSX, PPTX, HTML, and images to Markdown.</p>
<h2>Endpoints</h2>
<table>
<tr><th>Endpoint</th><th>Method</th><th>Description</th></tr>
<tr><td>/</td><td>GET</td><td>This documentation page</td></tr>
<tr><td>/health</td><td>GET</td><td>Health check with configuration</td></tr>
<tr><td>/marker</td><td>POST</td><td>Convert a file (sync, returns result)</td></tr>
<tr><td>/v1/conversions</td><td>POST</td><td>Convert a file (async-style, returns directly)</td></tr>
<tr><td>/v1/files/convert</td><td>POST</td><td>Convert a file via JSON body with base64</td></tr>
</table>

<h2>POST /marker <small>(multipart/form-data)</small></h2>
<pre><code>curl -X POST http://localhost:8000/marker \
  -F "file=@document.pdf" \
  -F "force_ocr=false" \
  -F "page_range=0,5-10" \
  -F "output_format=markdown"</code></pre>
<p>Returns the converted content directly as a file download.</p>

<h2>POST /marker <small>(application/json)</small></h2>
<pre><code>curl -X POST http://localhost:8000/marker \
  -H "Content-Type: application/json" \
  -d '{
    "file_b64": "base64-encoded-file-content",
    "filename": "document.pdf",
    "output_format": "markdown",
    "force_ocr": false
  }'</code></pre>

<h2>POST /v1/files/convert <small>(application/json)</small></h2>
<pre><code>curl -X POST http://localhost:8000/v1/files/convert \
  -H "Content-Type: application/json" \
  -d '{
    "file_b64": "base64-encoded-file-content",
    "filename": "document.pdf",
    "output_format": "json",
    "page_range": "0,5-10"
  }'</code></pre>

<h3>Options</h3>
<table>
<tr><th>Parameter</th><th>Type</th><th>Default</th><th>Description</th></tr>
<tr><td><b>file</b> / <b>file_b64</b></td><td>file / string</td><td><b>required</b></td><td>The file to convert</td></tr>
<tr><td>force_ocr</td><td>bool</td><td>false</td><td>Force OCR on all pages</td></tr>
<tr><td>paginate_output</td><td>bool</td><td>false</td><td>Separate pages with horizontal rules</td></tr>
<tr><td>output_format</td><td>string</td><td>markdown</td><td>markdown, json, html, chunks</td></tr>
<tr><td>page_range</td><td>string</td><td>all</td><td>Comma-separated page numbers/ranges: "0,5-10"</td></tr>
<tr><td>disable_image_extraction</td><td>bool</td><td>false</td><td>Disable extraction of embedded images</td></tr>
<tr><td>processors</td><td>string</td><td>auto</td><td>Comma-separated full module paths</td></tr>
<tr><td>config_json</td><td>string</td><td>none</td><td>Path to JSON file with additional config</td></tr>
<tr><td>converter_cls</td><td>string</td><td>auto-detected</td><td>Full module path of converter class</td></tr>
<tr><td>use_llm</td><td>bool</td><td>false</td><td>Use an LLM to improve accuracy (requires LLM service)</td></tr>
<tr><td>llm_service</td><td>string</td><td>marker.services.<br>ollama.OllamaService</td><td>LLM service class path: gemini, vertex, claude, openai, azure_openai, ollama</td></tr>
<tr><td>block_correction_prompt</td><td>string</td><td>none</td><td>Custom prompt for LLM block correction</td></tr>
<tr><td>redo_inline_math</td><td>bool</td><td>false</td><td>Re-process inline math with LLM</td></tr>
<tr><td>strip_existing_ocr</td><td>bool</td><td>false</td><td>Remove all existing OCR text and re-OCR</td></tr>
<tr><td>debug</td><td>bool</td><td>false</td><td>Enable debug mode with additional logging</td></tr>
</table>

<h3>Supported Formats</h3>
<p><code>{formats}</code></p>

<h3>Environment Variables</h3>
<table>
<tr><th>Variable</th><th>Default</th><th>Description</th></tr>
<tr><td>OLLAMA_HOST</td><td>http://10.0.1.127:11434</td><td>Ollama instance for OCR fallback</td></tr>
<tr><td>DEESEEK_OCR_MODEL</td><td>deepseek-ocr</td><td>OCR model name in Ollama</td></tr>
<tr><td>AMD_COMPUTE</td><td>false</td><td>Enable AMD ROCm GPU compute (set to "true")</td></tr>
<tr><td>TORCH_DEVICE</td><td>auto</td><td>PyTorch device: rocm, cuda, cpu</td></tr>
<tr><td>MODEL_DTYPE</td><td>float32</td><td>Model dtype: float32, bfloat16</td></tr>
<tr><td>PORT</td><td>8000</td><td>Listening port</td></tr>
<tr><td>HOST</td><td>0.0.0.0</td><td>Listening host</td></tr>
<tr><td>LLM_SERVICE</td><td>marker.services.ollama.OllamaService</td><td>Default LLM service class for use_llm</td></tr>
<tr><td>USE_LLM</td><td>false</td><td>Default use_llm flag (true/false)</td></tr>
<tr><td>OPENAI_BASE_URL</td><td></td><td>Base URL for OpenAI-compatible LLM service</td></tr>
<tr><td>OPENAI_API_KEY</td><td></td><td>API key for OpenAI-compatible LLM service</td></tr>
<tr><td>OPENAI_MODEL</td><td></td><td>Model name for OpenAI-compatible LLM service</td></tr>
</table>
</body></html>"""

    @app.route("/")
    def docs():
        body = HTML_DOCS.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_model": DEESEEK_OCR_MODEL,
            "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,
        })

    @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 = {
                "page_range": request.form.get("page_range"),
                "paginate_output": request.form.get("paginate_output", "false").lower() == "true",
                "force_ocr": request.form.get("force_ocr", "false").lower() == "true",
                "output_format": fmt,
                "disable_image_extraction": request.form.get("disable_image_extraction", "false").lower() == "true",
                "processors": request.form.get("processors"),
                "config_json": request.form.get("config_json"),
                "converter_cls": request.form.get("converter_cls"),
                "use_llm": request.form.get("use_llm", "false").lower() == "true",
                "llm_service": request.form.get("llm_service"),
                "block_correction_prompt": request.form.get("block_correction_prompt"),
                "redo_inline_math": request.form.get("redo_inline_math", "false").lower() == "true",
                "strip_existing_ocr": request.form.get("strip_existing_ocr", "false").lower() == "true",
                "debug": request.form.get("debug", "false").lower() == "true",
            }
            result = convert_file_bytes(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 = {
                "page_range": data.get("page_range"),
                "paginate_output": data.get("paginate_output", False),
                "force_ocr": data.get("force_ocr", False),
                "output_format": fmt,
                "disable_image_extraction": data.get("disable_image_extraction", False),
                "processors": data.get("processors"),
                "config_json": data.get("config_json"),
                "converter_cls": data.get("converter_cls"),
                "use_llm": data.get("use_llm", False),
                "llm_service": data.get("llm_service"),
                "block_correction_prompt": data.get("block_correction_prompt"),
                "redo_inline_math": data.get("redo_inline_math", False),
                "strip_existing_ocr": data.get("strip_existing_ocr", False),
                "debug": data.get("debug", False),
            }
            result = convert_file_bytes(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():
        """Convert a file, returns the result directly (async-style naming for API compatibility)."""
        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 = {
                "page_range": data.get("page_range"),
                "paginate_output": data.get("paginate_output", False),
                "force_ocr": data.get("force_ocr", False),
                "output_format": fmt,
                "disable_image_extraction": data.get("disable_image_extraction", False),
                "processors": data.get("processors"),
                "config_json": data.get("config_json"),
                "converter_cls": data.get("converter_cls"),
                "use_llm": data.get("use_llm", False),
                "llm_service": data.get("llm_service"),
                "block_correction_prompt": data.get("block_correction_prompt"),
                "redo_inline_math": data.get("redo_inline_math", False),
                "strip_existing_ocr": data.get("strip_existing_ocr", False),
                "debug": data.get("debug", False),
            }
            result = convert_file_bytes(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 = {
                "page_range": request.form.get("page_range"),
                "paginate_output": request.form.get("paginate_output", "false").lower() == "true",
                "force_ocr": request.form.get("force_ocr", "false").lower() == "true",
                "output_format": fmt,
                "disable_image_extraction": request.form.get("disable_image_extraction", "false").lower() == "true",
                "processors": request.form.get("processors"),
                "config_json": request.form.get("config_json"),
                "converter_cls": request.form.get("converter_cls"),
                "use_llm": request.form.get("use_llm", "false").lower() == "true",
                "llm_service": request.form.get("llm_service"),
                "block_correction_prompt": request.form.get("block_correction_prompt"),
                "redo_inline_math": request.form.get("redo_inline_math", "false").lower() == "true",
                "strip_existing_ocr": request.form.get("strip_existing_ocr", "false").lower() == "true",
                "debug": request.form.get("debug", "false").lower() == "true",
            }
            result = convert_file_bytes(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():
        """Convert a file via JSON body with base64."""
        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 = {
                "page_range": request.form.get("page_range"),
                "paginate_output": request.form.get("paginate_output", False),
                "force_ocr": request.form.get("force_ocr", False),
                "output_format": fmt,
                "disable_image_extraction": request.form.get("disable_image_extraction", False),
                "processors": request.form.get("processors"),
                "config_json": request.form.get("config_json"),
                "converter_cls": request.form.get("converter_cls"),
                "use_llm": request.form.get("use_llm", "false").lower() == "true",
                "llm_service": request.form.get("llm_service"),
                "block_correction_prompt": request.form.get("block_correction_prompt"),
                "redo_inline_math": request.form.get("redo_inline_math", "false").lower() == "true",
                "strip_existing_ocr": request.form.get("strip_existing_ocr", "false").lower() == "true",
                "debug": request.form.get("debug", "false").lower() == "true",
            }
            result = convert_file_bytes(raw, filename, **opts)
        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 = {
                "page_range": data.get("page_range"),
                "paginate_output": data.get("paginate_output", False),
                "force_ocr": data.get("force_ocr", False),
                "output_format": fmt,
                "disable_image_extraction": data.get("disable_image_extraction", False),
                "processors": data.get("processors"),
                "config_json": data.get("config_json"),
                "converter_cls": data.get("converter_cls"),
                "use_llm": data.get("use_llm", False),
                "llm_service": data.get("llm_service"),
                "block_correction_prompt": data.get("block_correction_prompt"),
                "redo_inline_math": data.get("redo_inline_math", False),
                "strip_existing_ocr": data.get("strip_existing_ocr", False),
                "debug": data.get("debug", False),
            }
            result = convert_file_bytes(raw, data["filename"], **opts)
            fmt = fmt
        else:
            return jsonify({"error": "No file provided."}), 400

        if not result["success"]:
            return jsonify(result), 500

        _filename = data.get("filename") if request.is_json else filename
        return jsonify({
            "filename": _filename,
            "format": fmt,
            "output": result["output"],
            "images_b64": result.get("images_b64", {}),
            "metadata": result.get("metadata", {}),
        })

    return app


app_instance = create_app()

if __name__ == "__main__":
    _configure_env()
    _marker_dict = create_model_dict()  # pre-warm models

    port = int(os.environ.get("PORT", "8000"))
    host = os.environ.get("HOST", "0.0.0.0")
    print("=" * 60)
    print("marker-api starting")
    print(f"  OLLAMA_HOST       = {OLLAMA_HOST}")
    print(f"  OCR_MODEL         = {DEESEEK_OCR_MODEL}")
    print(f"  AMD_COMPUTE       = {AMD_COMPUTE}")
    print(f"  TORCH_DEVICE      = {TORCH_DEVICE or '(auto)'}")
    print(f"  MODEL_DTYPE       = {MODEL_DTYPE}")
    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)