from . import models
from .database import SessionLocal

# Load environment variables from the directory where this file resides
base_dir = os.path.dirname(os.path.abspath(__file__))
dotenv_path = os.path.join(base_dir, ".env")
load_dotenv(dotenv_path)

def extract_label_info(image_bytes: bytes, mode: str = "item"):
    """
    Orchestrates extraction across multiple AI providers.
    Modes: 'item' (full technical extraction), 'box' (container discovery)
    """
    db = SessionLocal()
    try:
        if mode == "box":
            prompt = """
            Identify the CONTAINER or BOX name from this image.
            Look for large, prominent, bold, or hand-written text that identifies a storage unit.
            Ignore small technical details, quantities, or fine print.
            
            Return ONLY a valid JSON object:
            {
              "box_label": "The identified container name",
              "name": "Same as box_label",
              "category": "Storage",
              "description": "Brief description if useful",
              "quantity": 1
            }
            """
        else:
            # Fetch custom prompt from DB
            setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first()
            if setting:
                prompt = setting.value
            else:
                # Fallback to a sensible default if DB is not ready
                prompt = "Extract technical specs. Return JSON with name, category, description, connector, size, color, part_number, ocr_text, quantity."

        # 1. Try Gemini
        result = gemini.extract(image_bytes, prompt)
        
        if result:
            # Map user-defined prompt keys to model fields if needed
            # User keys: Item, Type, Description, Category, Connector, Size, Color, PartNr, OCR
            mapping = {
                "Item": "name",
                "Type": "type",
                "Description": "description",
                "Category": "category",
                "Connector": "connector",
                "Size": "size",
                "Color": "color",
                "PartNr": "part_number",
                "OCR": "ocr_text"
            }
            
            final_result = {}
            for ai_key, model_key in mapping.items():
                if ai_key in result:
                    final_result[model_key] = result[ai_key]
                elif model_key in result: # Already mapped or using model keys
                    final_result[model_key] = result[model_key]
            
            # Ensure quantity and barcode are handled if returned or default
            final_result["quantity"] = result.get("quantity", 1)
            final_result["barcode"] = result.get("barcode", result.get("PartNr", result.get("part_number", "")))
            
            # Handle Box mode specifically
            if mode == "box":
                final_result["box_label"] = result.get("box_label", result.get("name", "Unknown Box"))
                final_result["name"] = final_result["box_label"]

            return final_result

    finally:
        db.close()
        
    # 2. Try Claude (Fallback) - Note: Mapping logic would need to be replicated here if enabled
    # For now, keeping it simple
    return {"error": "AI extraction failed or no data returned. Check your API key and Prompt."}
        
    # 2. Try Claude (Fallback)
    result = claude.extract(image_bytes, prompt)
    if result:
        return result
        
    return {"error": "All AI providers failed or no API keys configured. Check your .env file."}