# app/services/ai_ocr_service.py
import json
import httpx
import base64
from app.schemas.evidence import RegistrationDocumentExtracted

class AiOcrService:
    OLLAMA_URL = "http://service_finder_ollama:11434/api/generate" 
    MODEL_NAME = "llama3.2-vision"

    @classmethod
    async def extract_registration_data(cls, clean_image_bytes: bytes) -> RegistrationDocumentExtracted:
        base64_image = base64.b64encode(clean_image_bytes).decode('utf-8')

        prompt = """
        Te egy magyar hatósági okmány-szakértő AI vagy. A feladatod a mellékelt magyar forgalmi engedély (kép) összes adatának kinyerése.
        
        Keresd meg és olvasd le az adatokat az alábbi hatósági kódok alapján:
        - A: Rendszám (kötőjellel, pl: ABC-123 vagy AA-BB-123)
        - B: Első nyilvántartásba vétel dátuma (YYYY.MM.DD)
        - C.1.1: Családi név vagy cégnév
        - C.1.2: Utónév
        - C.1.3: Teljes lakcím (Irsz, Város, Utca, Házszám)
        - C.4: Jogosultság (a = tulajdonos, b = üzembentartó)
        - D.1: Gyártmány (pl. TOYOTA, VOLKSWAGEN)
        - D.2: Jármű típusa
        - D.3: Kereskedelmi leírás (pl. COROLLA, GOLF)
        - E: Alvázszám (pontosan 17 karakter)
        - G: Saját tömeg (kg)
        - F.1: Együttes tömeg (kg)
        - P.1: Hengerűrtartalom (cm3)
        - P.2: Teljesítmény (kW)
        - P.3: Hajtóanyag (pl. Benzin, Gázolaj, Elektromos)
        - P.5: Motorkód
        - V.9: Környezetvédelmi osztály kódja
        - R: Szín
        - S.1: Ülések száma
        - H: Műszaki érvényesség vége (YYYY.MM.DD)
        - Sebességváltó: Keresd a 0, 1, 2, 3 kódokat (0=mechanikus, 2=automata).

        VÁLASZ FORMÁTUMA: Kizárólag érvényes JSON. Ha egy adat nem olvasható, az értéke null legyen.
        """

        payload = {
            "model": cls.MODEL_NAME,
            "prompt": prompt,
            "images": [base64_image],
            "stream": False,
            "format": "json"
        }

        async with httpx.AsyncClient(timeout=90.0) as client:
            try:
                response = await client.post(cls.OLLAMA_URL, json=payload)
                response.raise_for_status()
                
                ai_response_text = response.json().get("response", "{}")
                data_dict = json.loads(ai_response_text)
                
                return RegistrationDocumentExtracted(**data_dict)

            except Exception as e:
                print(f"Robot 3 AI Hiba: {e}")
                raise ValueError(f"AI hiba az adatkivonás során: {str(e)}")