# /opt/docker/dev/service_finder/backend/app/workers/ocr/robot_1_ocr_processor.py
import asyncio
import os
import logging
import json
from PIL import Image
from sqlalchemy import select, update
from app.db.session import AsyncSessionLocal
from app.models import Document
from app.models.identity import User
from app.models.marketplace.organization import Organization
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()
                        
                    # 1. Kép kódolása Base64 formátumba (A Vision modellek ezt várják)
                    import base64
                    base64_image = base64.b64encode(image_bytes).decode('utf-8')

                    # 2. A MEGFELELŐ AI HÍVÁS (Vision modell hívása az AIService-en keresztül)
                    doc_type_str = doc.doc_type or "general_document"
                    ocr_result = await AIService.process_ocr_document(
                        doc_type=doc_type_str, 
                        base64_image=base64_image
                    )

                    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)

                        # TRUST MATCHING: Keresés a fleet.organizations táblában adószám alapján
                        verified_org_id = None
                        tax_number = ocr_result.get("tax_number")
                        if tax_number:
                            org_stmt = select(Organization.id).where(
                                Organization.tax_number == tax_number,
                                Organization.is_active == True,
                                Organization.is_deleted == False
                            )
                            org_result = await db.execute(org_stmt)
                            org = org_result.scalar_one_or_none()
                            if org:
                                verified_org_id = org
                                logger.info(f"✅ Trust Matching sikeres: {tax_number} → org_id {verified_org_id}")
                            else:
                                logger.info(f"ℹ️ Trust Matching: nincs egyezés adószámra: {tax_number}")
                        
                        # OCR adatok frissítése verified_org_id-vel
                        if isinstance(ocr_result, dict):
                            ocr_result["verified_org_id"] = verified_org_id
                        else:
                            # Ha az ocr_result nem dict (pl. string), konvertáljuk
                            try:
                                ocr_dict = json.loads(ocr_result) if isinstance(ocr_result, str) else {}
                                ocr_dict["verified_org_id"] = verified_org_id
                                ocr_result = ocr_dict
                            except:
                                ocr_result = {"raw": ocr_result, "verified_org_id": verified_org_id}

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