# app/services/ai_ocr_service.py
import json
import httpx
import base64
import logging
from typing import Dict, Any, Optional
from app.schemas.evidence import RegistrationDocumentExtracted

logger = logging.getLogger(__name__)

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

    @classmethod
    async def analyze_image(cls, image_bytes: bytes, prompt: str) -> Dict[str, Any]:
        """
        Általános képfeldolgozás Ollama Vision modellel.
        
        Args:
            image_bytes: A kép bájtjai
            prompt: A prompt szöveg, amit a modelnek küldünk
        
        Returns:
            Dict a válasz adataival (a 'response' mezőből parse-olt JSON)
        
        Raises:
            httpx.RequestError: Ha a hálózati kérés sikertelen
            json.JSONDecodeError: Ha a válasz nem érvényes JSON
            ValueError: Ha más hiba történik
        """
        base64_image = base64.b64encode(image_bytes).decode('utf-8')
        
        payload = {
            "model": cls.MODEL_NAME,
            "prompt": prompt,
            "images": [base64_image],
            "stream": False,
            "format": "json"
        }
        
        async with httpx.AsyncClient(timeout=cls.DEFAULT_TIMEOUT) as client:
            try:
                logger.info(f"Ollama API hívás: {cls.OLLAMA_URL}, model: {cls.MODEL_NAME}")
                response = await client.post(cls.OLLAMA_URL, json=payload)
                response.raise_for_status()
                
                result = response.json()
                ai_response_text = result.get("response", "{}")
                
                # Próbáljuk JSON-ként értelmezni a választ
                try:
                    parsed = json.loads(ai_response_text)
                except json.JSONDecodeError:
                    # Ha nem JSON, visszaadjuk szövegként
                    parsed = {"raw_response": ai_response_text}
                
                logger.info(f"Ollama válasz sikeresen feldolgozva")
                return parsed
                
            except httpx.TimeoutException:
                logger.error("Ollama API timeout")
                raise ValueError("Ollama API időtúllépés")
            except httpx.HTTPStatusError as e:
                logger.error(f"Ollama HTTP hiba: {e.response.status_code} - {e.response.text}")
                raise ValueError(f"Ollama HTTP hiba: {e.response.status_code}")
            except Exception as e:
                logger.error(f"Ollama API hiba: {e}")
                raise ValueError(f"AI hiba a képfeldolgozás során: {str(e)}")

    @classmethod
    async def extract_registration_data(cls, clean_image_bytes: bytes) -> RegistrationDocumentExtracted:
        """
        Speciális metódus magyar forgalmi engedély adatainak kinyerésére.
        A régi kompatibilitás miatt megtartva.
        """
        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=cls.DEFAULT_TIMEOUT) 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:
                logger.error(f"Robot 3 AI Hiba: {e}")
                raise ValueError(f"AI hiba az adatkivonás során: {str(e)}")