"""Marker/OCR integration for document processing."""
from __future__ import annotations

import json
import asyncio
import io
import re
from pathlib import Path
from typing import Any
import httpx

from app.config import get_settings
import app.db.database as _database_mod


class MarkerProcessor:
    """Handles OCR via Marker API with polygon output."""

    def __init__(self):
        self.settings = get_settings()

    async def process_pdf(self, file_data: bytes, filename: str) -> dict[str, Any]:
        """Send PDF to Marker API for OCR with polygon extraction."""
        async with httpx.AsyncClient(timeout=300) as client:
            resp = await client.post(
                f"{self.settings.marker_api_url}/convert",
                files={"file": (filename, file_data, "application/pdf")},
            )
            resp.raise_for_status()
            return resp.json()

    async def process_pdf_url(self, pdf_url: str, filename: str) -> dict[str, Any]:
        """Process PDF from URL via Marker API."""
        async with httpx.AsyncClient(timeout=300) as client:
            resp = await client.post(
                f"{self.settings.marker_api_url}/convert",
                json={"pdf_url": pdf_url},
            )
            resp.raise_for_status()
            return resp.json()

    async def parse_marker_json(self, marker_output: dict, doc_id: str) -> list[dict]:
        """Parse Marker JSON output into vectorized chunks with polygon data."""
        pages = marker_output.get("pages", [])
        chunks = []
        
        for page in pages:
            page_num = page.get("meta", {}).get("page_num", page.get("page", 0))
            text_lines = page.get("text_lines", [])
            
            for block_idx, tl in enumerate(text_lines):
                content = tl.get("text", "")
                polygon = tl.get("bbox") or tl.get("polygon")
                block_type = tl.get("type", "text")
                
                if not content or not isinstance(content, str) or not content.strip():
                    continue
                
                chunks.append({
                    "content": content.strip(),
                    "page_num": page_num,
                    "block_index": block_idx,
                    "polygon": polygon,
                    "chunk_type": block_type,
                })

        if chunks:
            await _database_mod.db.batch_chunk(doc_id, chunks)
        
        return chunks

    async def process_document_file(
        self, file_data: bytes, filename: str, doc_id: str
    ) -> dict[str, Any]:
        """Process document file and return structured result."""
        ext = Path(filename).suffix.lower()
        
        if ext == ".pdf":
            result = await self.process_pdf(file_data, filename)
            pages_data = result.get("pages", [])
            
            chunks = await self.parse_marker_json({"pages": pages_data}, doc_id)
            return {
                "success": result.get("success", True),
                "doc_id": doc_id,
                "page_count": result.get("page_count", len(pages_data)),
                "chunks": len(chunks),
                "ocr_model": result.get("ocr_model", "deepseek-ocr"),
            }
        
        elif ext in (".txt", ".pdf", ".md"):
            text = file_data.decode("utf-8", errors="replace")
            paragraphs = re.split(r'\n\s*\n', text)
            chunks = []
            for i, para in enumerate(paragraphs):
                if len(para.strip()) > 20:
                    chunks.append({
                        "content": para.strip(),
                        "page_num": 0,
                        "block_index": i,
                        "polygon": None,
                        "chunk_type": "text",
                    })
            
            if chunks:
                await _database_mod.db.batch_chunk(doc_id, chunks)
            
            return {"success": True, "doc_id": doc_id, "chunks": len(chunks), "page_count": 1}

        return {"success": False, "error": f"Unsupported file type: {ext}"}
