"""
PaddleOCR-VL-1.5 Bridge Server (HF Spaces Edition)
====================================================
Deploys on Hugging Face Spaces as a FastAPI app.
Connects to vLLM Docker running on your GPU server.

Architecture:
    Gradio App (another HF Space or any client)
         |
    This HF Space (Bridge, port 7860)
         |
    Your GPU Server (vLLM Docker, 117.54.141.62:8000)

HF Space Settings → Variables and secrets:
    VLLM_SERVER_URL = http://117.54.141.62:8000/v1
    API_KEY = (optional, for auth)
"""

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

import uvicorn
from fastapi import FastAPI, File, Header, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from openai import OpenAI

# =============================================================================
# Configuration
# =============================================================================
VLLM_SERVER_URL = os.environ.get("VLLM_SERVER_URL", "http://117.54.141.62:8000/v1")
VLLM_MODEL_NAME = os.environ.get("VLLM_MODEL_NAME", "PaddleOCR-VL-1.5-0.9B")
BRIDGE_PORT = int(os.environ.get("PORT", "7860"))
API_KEY = os.environ.get("API_KEY", "")
# Public base URL for serving static files (auto-detect from HF Space)
SPACE_HOST = os.environ.get("SPACE_HOST", "")
if SPACE_HOST:
    PUBLIC_BASE_URL = f"https://{SPACE_HOST}"
else:
    PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", f"http://localhost:{BRIDGE_PORT}")

# Directory to store and serve output images
STATIC_DIR = "/tmp/ocr_outputs"
os.makedirs(STATIC_DIR, exist_ok=True)

# =============================================================================
# Initialize OpenAI client
# =============================================================================
openai_client = OpenAI(
    api_key="EMPTY",
    base_url=VLLM_SERVER_URL,
    timeout=600
)

# =============================================================================
# PaddleOCR pipeline
# =============================================================================
pipeline = None


def get_pipeline():
    """Lazy-load the PaddleOCR pipeline."""
    global pipeline
    if pipeline is None:
        from paddleocr import PaddleOCRVL
        pipeline = PaddleOCRVL(
            vl_rec_backend="vllm-server",
            vl_rec_server_url=VLLM_SERVER_URL
        )
    return pipeline


# =============================================================================
# FastAPI App
# =============================================================================
app = FastAPI(
    title="PaddleOCR-VL-1.5 Bridge API",
    description="Full document parsing API — bridge between Gradio UI and vLLM server",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve static files (output images)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


# =============================================================================
# Auth
# =============================================================================
def verify_auth(authorization: Optional[str] = None):
    if API_KEY and API_KEY.strip():
        if not authorization or authorization != f"Bearer {API_KEY}":
            raise HTTPException(status_code=401, detail="Unauthorized")


# =============================================================================
# Helpers
# =============================================================================
TASK_PROMPTS = {
    "ocr": "OCR:",
    "formula": "Formula Recognition:",
    "table": "Table Recognition:",
    "chart": "Chart Recognition:",
    "spotting": "Spotting:",
    "seal": "Seal Recognition:",
}

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}


def save_temp_image(file_data: str) -> str:
    """Save base64 or URL image to temp file."""
    if file_data.startswith(("http://", "https://")):
        import requests as req
        resp = req.get(file_data, timeout=120)
        resp.raise_for_status()
        content = resp.content
        ct = resp.headers.get("content-type", "image/png")
        ext = ".png"
        if "jpeg" in ct or "jpg" in ct:
            ext = ".jpg"
        elif "webp" in ct:
            ext = ".webp"
        elif "bmp" in ct:
            ext = ".bmp"
    else:
        content = base64.b64decode(file_data)
        ext = ".png"

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
    tmp.write(content)
    tmp.close()
    return tmp.name


def collect_output_images(output_dir: str, request_id: str) -> Dict[str, str]:
    """
    Find all image files in the output directory,
    copy them to the static dir, and return a dict of {name: public_url}.
    """
    output_images = {}
    if not os.path.exists(output_dir):
        return output_images

    # Create a subdirectory for this request
    static_subdir = os.path.join(STATIC_DIR, request_id)
    os.makedirs(static_subdir, exist_ok=True)

    for root, dirs, files in os.walk(output_dir):
        for filename in files:
            ext = os.path.splitext(filename)[1].lower()
            if ext in IMAGE_EXTENSIONS:
                src_path = os.path.join(root, filename)
                dst_path = os.path.join(static_subdir, filename)
                shutil.copy2(src_path, dst_path)
                public_url = f"{PUBLIC_BASE_URL}/static/{request_id}/{filename}"
                output_images[filename] = public_url

    return output_images


def element_level_recognition(file_data: str, prompt_label: str) -> Dict[str, Any]:
    """Element-level recognition via direct vLLM call."""
    if file_data.startswith(("http://", "https://")):
        image_url = file_data
    else:
        image_url = f"data:image/png;base64,{file_data}"

    task_prompt = TASK_PROMPTS.get(prompt_label, "OCR:")

    response = openai_client.chat.completions.create(
        model=VLLM_MODEL_NAME,
        messages=[{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": image_url}},
                {"type": "text", "text": task_prompt}
            ]
        }],
        temperature=0.0
    )

    result_text = response.choices[0].message.content

    return {
        "errorCode": 0,
        "result": {
            "layoutParsingResults": [{
                "markdown": {"text": result_text, "images": {}},
                "outputImages": {},
                "prunedResult": {
                    "spotting_res": _parse_spotting(result_text) if prompt_label == "spotting" else {}
                }
            }]
        }
    }


def full_document_parsing(file_data: str, use_chart_recognition: bool = False,
                          use_doc_unwarping: bool = True,
                          use_doc_orientation_classify: bool = True) -> Dict[str, Any]:
    """Full document parsing with layout detection + VLM recognition."""
    tmp_path = save_temp_image(file_data)
    request_id = str(uuid.uuid4())[:12]

    try:
        pipe = get_pipeline()
        output = pipe.predict(tmp_path)

        results = []
        for i, res in enumerate(output):
            output_dir = tempfile.mkdtemp()

            # Save all outputs (json, markdown, images)
            res.save_to_json(save_path=output_dir)
            res.save_to_markdown(save_path=output_dir)

            # Try to save visualization image
            try:
                res.save_to_img(save_path=output_dir)
            except Exception:
                pass

            # Read markdown
            md_text = ""
            md_files = [f for f in os.listdir(output_dir) if f.endswith(".md")]
            if md_files:
                with open(os.path.join(output_dir, md_files[0]), "r", encoding="utf-8") as f:
                    md_text = f.read()

            # Read JSON
            json_data = {}
            json_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]
            if json_files:
                with open(os.path.join(output_dir, json_files[0]), "r", encoding="utf-8") as f:
                    json_data = json.load(f)

            # Collect and serve output images
            page_request_id = f"{request_id}_page{i}"
            output_images = collect_output_images(output_dir, page_request_id)

            # Also check for images referenced in markdown
            md_images = {}
            for fname, url in output_images.items():
                # Replace local paths in markdown with public URLs
                md_text = md_text.replace(fname, url)
                md_images[fname] = url

            results.append({
                "markdown": {"text": md_text, "images": md_images},
                "outputImages": output_images,
                "jsonData": json_data
            })

        return {
            "errorCode": 0,
            "result": {
                "layoutParsingResults": results if results else [{
                    "markdown": {"text": "", "images": {}},
                    "outputImages": {}
                }]
            }
        }
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


def _parse_spotting(text: str) -> dict:
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        return {"raw_text": text}


# =============================================================================
# Endpoints
# =============================================================================

@app.get("/")
async def root():
    return {
        "service": "PaddleOCR-VL-1.5 Bridge API",
        "status": "running",
        "endpoints": ["/health", "/api/ocr", "/api/parse", "/api/parse/markdown", "/v1/chat/completions", "/docs"]
    }


@app.get("/health")
async def health():
    return {"status": "ok", "model": VLLM_MODEL_NAME, "vllm_url": VLLM_SERVER_URL}


@app.post("/api/ocr")
async def ocr_endpoint(request: Request, authorization: Optional[str] = Header(None)):
    """
    Main OCR endpoint — compatible with the Gradio app.

    Body:
    {
        "file": "base64_or_url",
        "useLayoutDetection": true/false,
        "promptLabel": "ocr|formula|table|chart|spotting|seal",
        "useChartRecognition": false,
        "useDocUnwarping": true,
        "useDocOrientationClassify": true
    }
    """
    verify_auth(authorization)

    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    file_data = body.get("file", "")
    if not file_data:
        raise HTTPException(status_code=400, detail="Missing 'file' field")

    use_layout = body.get("useLayoutDetection", False)
    prompt_label = body.get("promptLabel", "ocr")
    use_chart = body.get("useChartRecognition", False)
    use_unwarp = body.get("useDocUnwarping", True)
    use_orient = body.get("useDocOrientationClassify", True)

    try:
        if use_layout:
            return full_document_parsing(file_data, use_chart, use_unwarp, use_orient)
        else:
            return element_level_recognition(file_data, prompt_label)
    except Exception as e:
        traceback.print_exc()
        return {"errorCode": -1, "errorMsg": str(e)}


@app.post("/api/parse")
async def parse_file(
    file: UploadFile = File(...),
    use_layout_detection: bool = True,
    prompt_label: str = "ocr",
    authorization: Optional[str] = Header(None)
):
    """File upload endpoint."""
    verify_auth(authorization)
    content = await file.read()
    b64 = base64.b64encode(content).decode("utf-8")

    try:
        if use_layout_detection:
            return full_document_parsing(b64)
        else:
            return element_level_recognition(b64, prompt_label)
    except Exception as e:
        traceback.print_exc()
        return {"errorCode": -1, "errorMsg": str(e)}


@app.post("/api/parse/markdown")
async def parse_to_markdown(
    file: UploadFile = File(...),
    authorization: Optional[str] = Header(None)
):
    """Returns just markdown text."""
    verify_auth(authorization)
    content = await file.read()
    b64 = base64.b64encode(content).decode("utf-8")

    try:
        result = full_document_parsing(b64)
        pages = result.get("result", {}).get("layoutParsingResults", [])
        markdown_parts = [p.get("markdown", {}).get("text", "") for p in pages if p.get("markdown", {}).get("text")]
        return {
            "status": "ok",
            "markdown": "\n\n---\n\n".join(markdown_parts),
            "page_count": len(pages)
        }
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/v1/chat/completions")
async def proxy_chat_completions(request: Request, authorization: Optional[str] = Header(None)):
    """Proxy to vLLM for direct OpenAI-compatible calls."""
    verify_auth(authorization)

    import httpx
    body = await request.json()

    async with httpx.AsyncClient(timeout=600) as client:
        resp = await client.post(
            f"{VLLM_SERVER_URL}/chat/completions",
            json=body,
            headers={"Content-Type": "application/json"}
        )
        return resp.json()


# =============================================================================
# Entry point
# =============================================================================
if __name__ == "__main__":
    print(f"""
╔══════════════════════════════════════════════════════════════╗
║     PaddleOCR-VL-1.5 Bridge Server (HF Spaces)             ║
╠══════════════════════════════════════════════════════════════╣
║  Bridge API:   http://0.0.0.0:{BRIDGE_PORT}                          ║
║  vLLM backend: {VLLM_SERVER_URL:<44s}║
║  Model:        {VLLM_MODEL_NAME:<44s}║
║  Auth:         {"ENABLED" if API_KEY else "DISABLED":<44s}║
║  Static URL:   {PUBLIC_BASE_URL:<44s}║
╠══════════════════════════════════════════════════════════════╣
║  Endpoints:                                                  ║
║    GET  /health              - Health check                  ║
║    GET  /docs                - Swagger UI                    ║
║    POST /api/ocr             - Gradio-compatible API         ║
║    POST /api/parse           - File upload API               ║
║    POST /api/parse/markdown  - Simple markdown output        ║
║    POST /v1/chat/completions - vLLM proxy (OpenAI format)    ║
║    GET  /static/...          - Output images                 ║
╚══════════════════════════════════════════════════════════════╝
    """)
    uvicorn.run(app, host="0.0.0.0", port=BRIDGE_PORT)