"""
===============================================================================
FILE: services/office/document_extractor.py
PROJECT: SCUD Orion AI (Office Domain)
ROLE: Постраничное извлечение текста и OCR сканов документов через Vision LLM.
===============================================================================
"""

import os
import base64
import logging
from typing import List, Dict, Any, Optional

import fitz  # PyMuPDF
from modules.web_api.llm.core.ollama_client import call_ollama_chat

logger = logging.getLogger("OFFICE_EXTRACTOR")


def ocr_image_b64(image_b64: str, page_num: int = 1) -> str:
    """
    Распознает текст с одного растрового изображения через Qwen 2.5 VL.
    """
    system_prompt = (
        "Ты — высокоточный профессиональный модуль OCR для канцелярии и документооборота.\n"
        "Твоя задача — точно переписать весь текст с предоставленного изображения документа.\n"
        "ПРАВИЛА:\n"
        "1. Переписывай текст дословно, сохраняя структуру, заголовки, списки, нумерацию и таблицы.\n"
        "2. Запрещено добавлять вводные слова, приветствия, комментарии ('Спасибо за обращение', 'Вот текст' и т.д.).\n"
        "3. Выводи СТРОГО чистый распознанный текст документа."
    )
    user_message = {
        "role": "user",
        "content": f"Распознай весь печатный и рукописный текст страницы №{page_num} без пропусков.",
        "images": [image_b64]
    }
    try:
        res = call_ollama_chat(
            messages=[{"role": "system", "content": system_prompt}, user_message],
            is_vision=True,
            timeout=300  # 5 минут на тяжелые страницы
        )
        return (res.get("content") or "").strip()
    except Exception as e:
        logger.error(f"[Office OCR] Ошибка распознавания страницы {page_num}: {e}")
        return f"[Ошибка распознавания страницы {page_num}: {e}]"


def process_pdf_full(file_path: str, max_pages: Optional[int] = None) -> List[Dict[str, Any]]:
    """
    Обрабатывает PDF целиком:
    - Если есть качественный цифровой текст — мгновенно извлекает его со всех страниц.
    - Если страница является сканом/картинкой — рендерит ее в 200 DPI и прогоняет через Vision OCR.
    """
    results = []
    doc = fitz.open(file_path)
    total_pages = len(doc)
    limit = min(total_pages, max_pages) if max_pages else total_pages

    logger.info(f"[Office] Начало обработки PDF: {os.path.basename(file_path)} (всего страниц: {total_pages})")

    for idx in range(limit):
        page_num = idx + 1
        page = doc[idx]
        extracted_text = (page.get_text("text") or "").strip()

        # Если на странице есть хороший машинный текст (не скан)
        if len(extracted_text) > 80:
            logger.info(f"[Office] Страница {page_num}/{limit}: извлечен цифровой текст ({len(extracted_text)} симв.)")
            results.append({
                "page": page_num,
                "method": "DIGITAL_TEXT",
                "text": extracted_text
            })
        else:
            # Чистый скан — рендерим страницу в PNG и передаем в Vision LLM
            logger.info(f"[Office] Страница {page_num}/{limit}: распознавание скана через Ollama Vision...")
            pix = page.get_pixmap(dpi=200)
            img_b64 = base64.b64encode(pix.tobytes("png")).decode("utf-8")
            
            ocr_text = ocr_image_b64(img_b64, page_num=page_num)
            results.append({
                "page": page_num,
                "method": "VISION_OCR",
                "text": ocr_text
            })

    doc.close()
    return results