# /opt/docker/dev/service_finder/backend/app/workers/vehicle/vehicle_robot_3_alchemist_pro.py
"""
Robot 3: Alchemist Pro - AI Szintézis és Kapuőr
Javítások: 
- Batch Size: 3 (Stabilitás a 14b modellhez)
- Szigorú Gatekeeper (Arany státusz ellenőrzés)
- Adatmegőrzés: Az AI nem bírálja felül a szótár alapú RDW adatokat (kW/ccm).
"""

import asyncio
import logging
import sys
import json
import re
from sqlalchemy import text, update, func
from sqlalchemy.ext.asyncio import AsyncSession
import httpx
from app.database import AsyncSessionLocal
from app.models import VehicleModelDefinition

logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s [%(levelname)s] R3-Alchemist-Pro: %(message)s', 
    stream=sys.stdout
)
logger = logging.getLogger("Robot-3-Alchemist-Pro")

OLLAMA_URL = "http://sf_ollama:11434/api/generate"
OLLAMA_MODEL = "qwen2.5-coder:14b" # A 14b paraméteres modell az agy
MAX_ATTEMPTS = 3
TIMEOUT_SECONDS = 45 # Megemelt timeout a 14b modell lassabb válaszideje miatt
BATCH_SIZE = 4 # Maximum 4 párhuzamos AI hívás a CPU fagyás elkerülésére

class AlchemistPro:
    def __init__(self):
        self.client = httpx.AsyncClient(timeout=TIMEOUT_SECONDS)

    async def close(self):
        await self.client.aclose()

    async def fetch_vehicle_batch_for_processing(self, db: AsyncSession):
        """Kiválasztja azokat a járműveket, ahol a 2.1-es robot végzett, de még nem 'Arany'."""
        query = text("""
            SELECT id, make, marketing_name, power_kw, engine_capacity,
                   fuel_type, raw_api_data, raw_search_context, attempts,
                   vehicle_class, trim_level, transmission_type, body_type
            FROM vehicle.vehicle_model_definitions
            WHERE status = 'awaiting_ai_synthesis'
              AND attempts < :max_attempts
              AND is_manual = FALSE
            ORDER BY priority_score DESC NULLS LAST, id ASC
            FOR UPDATE SKIP LOCKED
            LIMIT :batch_size
        """)
        result = await db.execute(query, {"max_attempts": MAX_ATTEMPTS, "batch_size": BATCH_SIZE})
        rows = result.fetchall()
        
        vehicles = []
        for row in rows:
            vehicles.append({
                "id": row[0], "make": row[1], "marketing_name": row[2],
                "power_kw": row[3] or 0, "engine_capacity": row[4] or 0,
                "fuel_type": row[5] or "Unknown", "raw_api_data": row[6] or {},
                "raw_search_context": row[7] or "", "attempts": row[8] or 0,
                "vehicle_class": row[9], "trim_level": row[10],
                "transmission_type": row[11], "body_type": row[12]
            })
        return vehicles

    def build_prompt(self, vehicle_data: dict) -> str:
        """Megfogalmazza a feladatot az AI számára a 14b modell erejét kihasználva."""
        make = vehicle_data["make"]
        model = vehicle_data["marketing_name"]
        
        # Rövidítjük a kontextust, hogy beleférjen a kontextus ablakba
        raw_api = json.dumps(vehicle_data["raw_api_data"], ensure_ascii=False)[:1000]
        raw_context = (vehicle_data["raw_search_context"] or "")[:2000]

        prompt = f"""
        Analyze the vehicle data and return missing information in valid JSON format.
        Vehicle: {make} {model}
        Current Specs:
        - Power: {vehicle_data['power_kw']} kW (0 means missing)
        - Engine: {vehicle_data['engine_capacity']} ccm (0 means missing)
        - Fuel: {vehicle_data['fuel_type']}
        
        Context Data: {raw_api}
        Search Snippets: {raw_context}

        INSTRUCTIONS:
        1. Identify trim_level (e.g., GTI, AMG, Highline, Titanium).
        2. Identify transmission (MANUAL, AUTOMATIC, CVT, DCT).
        3. Identify body_type (SEDAN, SUV, HATCHBACK, ESTATE, COUPE).
        4. If Power is 0, estimate it based on the engine size and fuel in context.
        5. If Engine is 0, estimate it based on model name.

        Return ONLY a JSON object:
        {{
            "trim_level": "string",
            "transmission": "string",
            "body_type": "string",
            "estimated_kw": integer_or_null,
            "estimated_ccm": integer_or_null
        }}
        """
        return prompt.strip()

    async def call_ollama(self, prompt: str) -> dict:
        """Kommunikáció az Ollama szerverrel."""
        payload = {
            "model": OLLAMA_MODEL,
            "prompt": prompt,
            "format": "json",
            "stream": False,
            "options": {"temperature": 0.1, "top_p": 0.9, "num_ctx": 4096}
        }
        try:
            response = await self.client.post(OLLAMA_URL, json=payload)
            response.raise_for_status()
            data = response.json()
            return json.loads(data.get("response", "{}"))
        except Exception as e:
            raise ValueError(f"Ollama hiba: {str(e)}")

    def merge_vehicle_data(self, vehicle: dict, ai_result: dict) -> dict:
        """Összefésüli a meglévő adatokat az AI eredményeivel, prioritást adva a meglévőnek."""
        merged = vehicle.copy()
        
        # A szöveges mezőket frissítjük, ha az AI talált jobbat
        for field, ai_key in [("trim_level", "trim_level"), ("transmission_type", "transmission"), ("body_type", "body_type")]:
            if not merged.get(field) and ai_result.get(ai_key):
                merged[field] = str(ai_result[ai_key]).upper() if field != "trim_level" else ai_result[ai_key]

        # MATEK VÉDELEM: Csak akkor írjuk be az AI becslését, ha a 2.1-es robot nem talált adatot (még mindig 0)
        if merged["power_kw"] == 0 and ai_result.get("estimated_kw"):
            merged["power_kw"] = int(ai_result["estimated_kw"])
        
        if merged["engine_capacity"] == 0 and ai_result.get("estimated_ccm"):
            merged["engine_capacity"] = int(ai_result["estimated_ccm"])
            
        return merged

    async def update_vehicle_record(self, db: AsyncSession, vehicle_id: int, merged_data: dict):
        """Végrehajtja a mentést és a Kapuőr logikát."""
        kw = merged_data.get("power_kw", 0)
        ccm = merged_data.get("engine_capacity", 0)
        fuel = str(merged_data.get("fuel_type", "")).lower()
        v_class = str(merged_data.get("vehicle_class", "")).lower()

        # Kapuőr szabályok
        is_electric = any(x in fuel for x in ['electr', 'elektri', 'hydrogen'])
        is_trailer = 'trailer' in v_class
        
        is_gold = False
        if is_trailer: is_gold = True
        elif is_electric: is_gold = kw > 0
        else: is_gold = (kw > 0 and ccm > 0)

        if is_gold:
            new_status = "gold_enriched"
            new_attempts = 0
            msg = "✨ ARANY"
        else:
            new_attempts = merged_data["attempts"] + 1
            new_status = "manual_review_needed" if new_attempts >= MAX_ATTEMPTS else "unverified"
            msg = "🔄 VISSZADOBVA"

        update_values = {
            "trim_level": merged_data.get("trim_level"),
            "transmission_type": merged_data.get("transmission_type"),
            "body_type": merged_data.get("body_type"),
            "power_kw": kw,
            "engine_capacity": ccm,
            "status": new_status,
            "attempts": new_attempts,
            "updated_at": func.now()
        }
        
        stmt = update(VehicleModelDefinition).where(VehicleModelDefinition.id == vehicle_id).values(**update_values)
        await db.execute(stmt)
        logger.info(f"{msg}: {merged_data['make']} {merged_data['marketing_name']} (Státusz: {new_status})")

    async def process_ai_task(self, vehicle: dict):
        """AI feldolgozás párhuzamosítható része."""
        try:
            prompt = self.build_prompt(vehicle)
            ai_result = await self.call_ollama(prompt)
            return vehicle, ai_result, None
        except Exception as e:
            return vehicle, None, e

    async def process_batch(self, db: AsyncSession, vehicles: list):
        """Batch feldolgozás: Szekvenciális AI feldolgozás a VRAM korlátok miatt."""
        results = []
        
        # 1. AI kérések szekvenciálisan (egy jármű után a másik)
        for vehicle in vehicles:
            try:
                vehicle_result = await self.process_ai_task(vehicle)
                results.append(vehicle_result)
            except Exception as e:
                logger.error(f"Hiba {vehicle['id']} AI feldolgozás közben: {e}")
                results.append((vehicle, None, e))
        
        # 2. Mentés szekvenciálisan a DB lakatok elkerülésére
        for vehicle, ai_result, error in results:
            if error:
                logger.error(f"Hiba {vehicle['id']}: {error}")
                # Hiba esetén növeljük a próbálkozások számát
                stmt = update(VehicleModelDefinition).where(VehicleModelDefinition.id == vehicle['id']).values(
                    attempts=vehicle['attempts'] + 1,
                    updated_at=func.now()
                )
                await db.execute(stmt)
            else:
                merged = self.merge_vehicle_data(vehicle, ai_result)
                await self.update_vehicle_record(db, vehicle['id'], merged)
        
        await db.commit()

    async def run(self):
        logger.info(f"🚀 Robot 3 indítva. Modell: {OLLAMA_MODEL}, Batch: {BATCH_SIZE}")
        while True:
            try:
                async with AsyncSessionLocal() as db:
                    vehicles = await self.fetch_vehicle_batch_for_processing(db)
                    if vehicles:
                        logger.info(f"📦 Feldolgozás: {len(vehicles)} jármű...")
                        await self.process_batch(db, vehicles)
                        await asyncio.sleep(1)
                    else:
                        await asyncio.sleep(2)
            except Exception as e:
                logger.error(f"Főciklus hiba: {e}")
                await asyncio.sleep(5)

if __name__ == "__main__":
    robot = AlchemistPro()
    asyncio.run(robot.run())