# /opt/docker/dev/service_finder/backend/app/workers/ocr_robot.py
import asyncio
import os
import logging
from PIL import Image
from sqlalchemy import select, update
from app.db.session import AsyncSessionLocal
from app.models.document import Document
from app.models.identity import User
from app.services.ai_service import AIService
from app.core.config import settings

# Logolás beállítása
logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(name)s: %(message)s')
logger = logging.getLogger("Robot-OCR-V3")

class OCRRobot:
    """ 
    Robot 3: Dokumentum elemző és adatkinyerő. 
    Kizárólag a Premium és VIP előfizetők dokumentumait dolgozza fel automatikusan.
    """

    @staticmethod
    def _sync_resize_and_save(source: str, target: str):
        """ Kép optimalizálása (szinkron végrehajtás a Pillow miatt). """
        with Image.open(source) as img:
            # Konvertálás RGB-be (PNG/RGBA -> JPEG támogatás miatt)
            rgb_img = img.convert('RGB')
            # Max szélesség 1600px az MB 2.0 Vault szabályai szerint
            if rgb_img.width > 1600:
                ratio = 1600 / float(rgb_img.width)
                new_height = int(float(rgb_img.height) * float(ratio))
                rgb_img = rgb_img.resize((1600, new_height), Image.Resampling.LANCZOS)
            
            rgb_img.save(target, "JPEG", quality=85, optimize=True)

    @classmethod
    async def process_queue(cls):
        """ A várólista feldolgozása. """
        async with AsyncSessionLocal() as db:
            # 1. LOGIKA: Feladatok lekérése (Pending + Premium jogosultság)
            # A 'SKIP LOCKED' biztosítja, hogy több robot ne akadjon össze
            stmt = select(Document, User).join(User, Document.parent_id == User.scope_id).where(
                Document.status == "pending_ocr",
                User.subscription_plan.in_(["PREMIUM_PLUS", "VIP_PLUS", "PREMIUM", "VIP"])
            ).limit(5)
            
            res = await db.execute(stmt)
            tasks = res.all()

            if not tasks:
                return

            for doc, user in tasks:
                try:
                    logger.info(f"📸 OCR megkezdése: {doc.original_name} (Szervezet: {user.scope_id})")
                    
                    # Státusz zárolása
                    doc.status = "processing"
                    await db.commit()

                    # 2. LOGIKA: AI OCR hívás az AIService-en keresztül
                    # Itt feltételezzük, hogy a Document modellben tároljuk a temp_path-t
                    if not doc.file_hash: # Biztonsági check
                         raise ValueError("Hiányzó fájl hivatkozás.")

                    temp_path = f"/app/temp/uploads/{doc.file_hash}"
                    
                    if not os.path.exists(temp_path):
                        raise FileNotFoundError(f"A forrásfájl nem található: {temp_path}")

                    with open(temp_path, "rb") as f:
                        image_bytes = f.read()

                    # AI felismerés (pl. Llama-Vision vagy GPT-4o)
                    ocr_result = await AIService.get_clean_vehicle_data(
                        make="OCR_SCAN", 
                        raw_model=doc.parent_type, 
                        v_type="document", 
                        sources={"image_data": "raw_scan"}
                    )

                    if ocr_result:
                        # 3. LOGIKA: Vault mentés (NAS izoláció)
                        target_dir = os.path.join(settings.NAS_STORAGE_PATH, user.folder_slug or "common", "vault")
                        os.makedirs(target_dir, exist_ok=True)
                        
                        final_filename = f"{doc.id}.jpg"
                        final_path = os.path.join(target_dir, final_filename)

                        # Kép feldolgozása külön szálon, hogy ne blokkolja az Async-et
                        loop = asyncio.get_event_loop()
                        await loop.run_in_executor(None, cls._sync_resize_and_save, temp_path, final_path)

                        # 4. LOGIKA: Adatbázis frissítés (Gold Data előkészítés)
                        doc.ocr_data = ocr_result
                        doc.status = "processed"
                        doc.file_size = os.path.getsize(final_path)
                        
                        # Ideiglenes fájl takarítása
                        os.remove(temp_path)
                        logger.info(f"✅ Dokumentum sikeresen archiválva: {final_filename}")
                    else:
                        doc.status = "failed"
                        doc.error_log = "AI returned empty result"

                    await db.commit()

                except Exception as e:
                    logger.error(f"❌ OCR Kritikus Hiba ({doc.id}): {str(e)}")
                    await db.rollback()
                    # Hibás státusz mentése
                    async with AsyncSessionLocal() as error_db:
                        await error_db.execute(
                            update(Document).where(Document.id == doc.id).values(
                                status="failed", 
                                error_log=str(e)
                            )
                        )
                        await error_db.commit()

    @classmethod
    async def run(cls):
        """ Folyamatos futtatás (Service mode). """
        logger.info("🤖 Robot 3 (OCR) ONLINE - Figyeli a prémium dokumentumokat")
        while True:
            await cls.process_queue()
            await asyncio.sleep(15) # 15 másodpercenkénti ellenőrzés

if __name__ == "__main__":
    asyncio.run(OCRRobot.run())