import os, base64, json, re
import httpx
from fastapi import APIRouter, File, UploadFile, HTTPException, Body
from langchain_google_genai import ChatGoogleGenerativeAI
from dotenv import load_dotenv
from typing import Dict, List, Any
from pathlib import Path

router = APIRouter(
    prefix="/ai",
    tags=["AI"]
)

BASE_DIR = Path(__file__).resolve().parent.parent.parent
env_path = BASE_DIR / ".env.local"
if env_path.exists():
    load_dotenv(dotenv_path=env_path)
else:
    print(f"Warning: .env.local file not found at {env_path}")
    print("Please copy .env.template to .env.local and add your GEMINI_API_KEY")

api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    print("Warning: GEMINI_API_KEY not set. AI features will be disabled.")
    print("Please set your GEMINI_API_KEY in .env.local file")
    api_key = None

llm = None
if api_key:
    try:
        llm = ChatGoogleGenerativeAI(
            model="gemini-2.0-flash",
            google_api_key=api_key
        )
    except Exception as e:
        print(f"Error initializing LLM: {e}")
        llm = None

def encode_image(image_content: bytes) -> str:
    return base64.b64encode(image_content).decode()

def estimate_manpower(weight: int) -> int:
    """
    Suggest manpower (including 1 driver).
    Rule of thumb:
    - ≤1 ton → 1 driver + 1 mover = 2
    - ≤3 ton → 1 driver + 2 movers = 3
    - ≤5 ton → 1 driver + 3 movers = 4
    """
    if weight <= 1000:
        return 2
    elif weight <= 3000:
        return 3
    else:
        return 4

def get_smart_alternatives(item_analysis: dict, lorry_data: list, photo_count: int = 0) -> list:
    """
    AI Learning: Smart filtering of alternatives based on detected items and photo count.
    Enhanced to provide realistic recommendations based on actual item volume and weight.
    """
    # Extract lorry information
    lorry_map = {}
    for lorry in lorry_data:
        lorry_map[lorry.get("id")] = {
            "lorry_size": lorry.get("lorry_size"),
            "weight_kg": lorry.get("weight_kg", 0)
        }
    
    # Enhanced filtering based on actual item analysis
    confidence = item_analysis.get("confidence", "low")
    has_only_small_items = item_analysis.get("has_only_small_items", False)
    small_item_count = item_analysis.get("small_item_count", 0)
    large_item_count = item_analysis.get("large_item_count", 0)
    
    # Calculate total item complexity
    total_items = small_item_count + large_item_count
    
    # Realistic recommendations based on actual items
    if has_only_small_items and small_item_count <= 3:
        # Very small moves: Van or Small lorry
        return [4, 1]  # Van first, then Small (1-Ton)
    
    elif has_only_small_items and small_item_count <= 8:
        # Small moves with multiple items: Small lorry or Van
        return [1, 4]  # Small (1-Ton) first, then Van
    
    elif has_only_small_items and small_item_count > 8:
        # Many small items: Medium lorry might be needed
        return [1, 2]  # Small (1-Ton) first, then Medium (3-Ton)
    
    elif large_item_count > 0:
        # Has furniture/appliances: Use appropriate sizing
        if large_item_count <= 2 and small_item_count <= 5:
            # Few furniture pieces: Medium lorry
            return [2, 4]  # Medium (3-Ton) first, then Van
        elif large_item_count <= 4 and small_item_count <= 10:
            # Moderate furniture: Medium lorry
            return [2, 3]  # Medium (3-Ton) first, then Large (5-Ton)
        else:
            # Many furniture pieces: Large lorry
            return [3, 2]  # Large (5-Ton) first, then Medium (3-Ton)
    
    # Default fallback: Provide reasonable alternatives
    else:
        return [2, 1]  # Medium (3-Ton) first, then Small (1-Ton)

def analyze_item_types(ai_response: str) -> dict:
    """
    AI Learning: Analyze AI response to detect item types and determine if only small items.
    Enhanced box detection with careful analysis.
    """
    ai_response_lower = ai_response.lower()
    
    # Enhanced keywords for small items with careful box detection
    small_item_keywords = [
        # Box variations - carefully detected
        'box', 'boxes', 'cardboard box', 'cardboard boxes', 'moving box', 'moving boxes',
        'storage box', 'storage boxes', 'packing box', 'packing boxes',
        'carton', 'cartons', 'cardboard carton', 'cardboard cartons',
        'package', 'packages', 'wrapped package', 'wrapped packages',
        'parcel', 'parcels', 'small parcel', 'small parcels',
        
        # Bag variations
        'bag', 'bags', 'plastic bag', 'plastic bags', 'shopping bag', 'shopping bags',
        'grocery bag', 'grocery bags', 'cloth bag', 'cloth bags',
        
        # Luggage and travel items
        'suitcase', 'suitcases', 'luggage', 'travel bag', 'travel bags',
        'backpack', 'backpacks', 'duffel bag', 'duffel bags',
        
        # Small personal items
        'book', 'books', 'notebook', 'notebooks', 'magazine', 'magazines',
        'clothes', 'clothing', 'shirt', 'shirts', 'pants', 'trousers',
        'dress', 'dresses', 'jacket', 'jackets', 'sweater', 'sweaters',
        
        # Electronics and small devices
        'laptop', 'laptops', 'computer', 'computers', 'phone', 'phones',
        'tablet', 'tablets', 'camera', 'cameras', 'headphone', 'headphones',
        
        # Documents and papers
        'document', 'documents', 'paper', 'papers', 'file', 'files',
        'folder', 'folders', 'envelope', 'envelopes',
        
        # Small decorative items
        'decorative', 'decoration', 'decorations', 'ornament', 'ornaments',
        'personal items', 'accessories', 'jewelry', 'jewellery',
        'toy', 'toys', 'small toy', 'small toys',
        
        # Kitchen and bathroom small items
        'cup', 'cups', 'plate', 'plates', 'bowl', 'bowls', 'mug', 'mugs',
        'utensil', 'utensils', 'cutlery', 'small appliance', 'small appliances'
    ]
    
    # Enhanced keywords for large items with careful detection
    large_item_keywords = [
        # Furniture - carefully detected
        'furniture', 'piece of furniture', 'pieces of furniture',
        'sofa', 'sofas', 'couch', 'couches', 'settee', 'settees',
        'bed', 'beds', 'bedroom set', 'bedroom sets', 'mattress', 'mattresses',
        'table', 'tables', 'dining table', 'dining tables', 'coffee table', 'coffee tables',
        'chair', 'chairs', 'dining chair', 'dining chairs', 'armchair', 'armchairs',
        'desk', 'desks', 'office desk', 'office desks',
        
        # Large appliances
        'refrigerator', 'refrigerators', 'fridge', 'fridges', 'freezer', 'freezers',
        'washing machine', 'washing machines', 'dryer', 'dryers',
        'dishwasher', 'dishwashers', 'oven', 'ovens', 'stove', 'stoves',
        
        # Large storage and cabinets
        'cabinet', 'cabinets', 'wardrobe', 'wardrobes', 'dresser', 'dressers',
        'bookshelf', 'bookshelves', 'shelf', 'shelves', 'closet', 'closets',
        'chest of drawers', 'chests of drawers',
        
        # Heavy and large items
        'large', 'heavy', 'big', 'huge', 'massive', 'bulky',
        'piano', 'pianos', 'grand piano', 'upright piano',
        'treadmill', 'treadmills', 'exercise equipment', 'gym equipment',
        
        # Construction and outdoor items
        'ladder', 'ladders', 'tool', 'tools', 'workbench', 'workbenches',
        'garden furniture', 'patio furniture', 'outdoor furniture'
    ]
    
    # Enhanced detection with context analysis
    small_item_count = 0
    large_item_count = 0
    
    # Count small items with context checking
    for keyword in small_item_keywords:
        if keyword in ai_response_lower:
            # Check if it's not part of a larger item description
            keyword_index = ai_response_lower.find(keyword)
            before_context = ai_response_lower[max(0, keyword_index-20):keyword_index]
            after_context = ai_response_lower[keyword_index+len(keyword):min(len(ai_response_lower), keyword_index+len(keyword)+20)]
            
            # Avoid counting if it's part of "large box" or "big box" etc.
            if not any(word in before_context + after_context for word in ['large', 'big', 'huge', 'massive', 'bulky']):
                small_item_count += 1
    
    # Count large items with context checking
    for keyword in large_item_keywords:
        if keyword in ai_response_lower:
            large_item_count += 1
    
    # Enhanced confidence calculation
    confidence = "low"
    if small_item_count > 0 and large_item_count == 0:
        if small_item_count >= 5:
            confidence = "very high"
        elif small_item_count >= 3:
            confidence = "high"
        else:
            confidence = "medium"
    elif small_item_count > large_item_count:
        confidence = "medium"
    else:
        confidence = "low"
    
    return {
        "has_only_small_items": small_item_count > 0 and large_item_count == 0,
        "small_item_count": small_item_count,
        "large_item_count": large_item_count,
        "confidence": confidence,
        "analysis_details": {
            "small_items_detected": [kw for kw in small_item_keywords if kw in ai_response_lower],
            "large_items_detected": [kw for kw in large_item_keywords if kw in ai_response_lower],
            "context_analysis": "Enhanced detection with context checking applied"
        }
    }

@router.post("/lorry-estimator")
async def lorry_estimator(files: List[UploadFile] = File(...)) -> Dict[str, Any]:
    # Debug logging
    print(f"🔍 Received {len(files) if files else 0} files")
    if files:
        for i, file in enumerate(files):
            print(f"🔍 File {i}: {file.filename}, Content-Type: {file.content_type}, Size: {file.size}")
    
    if not llm:
        raise HTTPException(503, detail="AI service unavailable. Please configure GEMINI_API_KEY in .env.local file.")
    
    if not files:
        raise HTTPException(400, detail="No files uploaded.")

    # Fetch lorry list
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get("http://127.0.0.1:8000/lorries/")
            response.raise_for_status()
            lorry_data = response.json()
        except httpx.HTTPError as e:
            raise HTTPException(500, detail=f"Failed to fetch lorry list: {e}")

    if not lorry_data:
        raise HTTPException(500, detail="Lorry list is empty.")

    # Lorry mapping (example, replace with DB fetch)
    lorry_map = {
        1: {"lorry_size": "Small (1-Ton)", "weight_kg": 1000},
        2: {"lorry_size": "Medium (3-Ton)", "weight_kg": 3000},
        3: {"lorry_size": "Large (5-Ton)", "weight_kg": 5000},
        4: {"lorry_size": "Van (0.8-Ton)", "weight_kg": 800}, # Van for boxes and small items
    }

    # Lorries summary for prompt
    lorries_for_prompt = [
        {
            "id": item.get("id"),
            "lorry_size": item.get("lorry_size"),
            "length_m": item.get("length_m"),
            "width_m": item.get("width_m"),
            "height_m": item.get("height_m"),
            "weight_kg": item.get("weight_kg"),
        }
        for item in lorry_data
    ]
    lorry_info_str = json.dumps(lorries_for_prompt)

    # Ini bahagian AI analysis for each photo
    user_content = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": (
                        f"Available lorries: {lorry_info_str}\n\n"
                        "TASK: Analyze each uploaded image individually and provide a detailed breakdown.\n\n"
                        
                        "ANALYSIS FORMAT:\n"
                        "For each image, provide:\n"
                        "**Image 1:** [Brief description - Bedroom/Living Room/Kitchen/Empty Room/Outdoor Scene/etc.]\n"
                        "- Items detected: [List specific moving items OR 'None' if no relevant items]\n"
                        "- Key observations: [Size, condition, special notes OR 'Irrelevant for moving' if empty/outdoor]\n\n"
                        
                        "**Image 2:** [Brief description]\n"
                        "- Items detected: [List specific items OR 'None']\n"
                        "- Key observations: [Notable features OR 'Irrelevant for moving']\n\n"
                        
                        "(Continue for all images)\n\n"
                        
                        "IMPORTANT RULES FOR IMAGE CLASSIFICATION:\n"
                        "- If image shows empty rooms/spaces: Items detected = 'None', Observations = 'Empty space'\n"
                        "- If image shows outdoor scenes/landscapes: Items detected = 'None', Observations = 'Outdoor scene - not relevant for moving'\n"
                        "- If image shows only walls/ceilings/floors: Items detected = 'None', Observations = 'Interior view - no movable items'\n"
                        "- If image shows people without furniture: Items detected = 'None', Observations = 'People only - no moving items'\n"
                        "- Only count actual furniture, boxes, and movable household items\n\n"
                        
                        "**Overall Assessment:**\n"
                        "- Total furniture pieces: [count]\n"
                        "- Total boxes/containers: [estimated count and sizes]\n"
                        "- Heavy/special items: [list any]\n"
                        "- Recommended lorry: [lorry size] (ID: [number])\n"
                        "- Alternative options: [provide 2-3 alternative lorry IDs with brief reasoning]\n"
                        "- Reasoning: [brief explanation for main recommendation]\n"
                        "- Manpower needed: [number] movers (excluding driver)\n\n"
                        
                        "RULES:\n"
                        "- Be specific about what you see in each image\n"
                        "- Count items realistically\n"
                        "- Consider total volume for lorry recommendation\n"
                        "- Base manpower on item complexity and weight"
                    )
                }
            ]
        }
    ]

    # Add images to the prompt
    for file in files:
        if file.content_type and file.content_type.startswith("image/"):
            try:
                image_content = await file.read()
                encoded_image = encode_image(image_content)
                user_content[0]["content"].append({
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{encoded_image}"
                    }
                })
            except Exception as e:
                print(f"Error processing image {file.filename}: {e}")
                continue

    try:
        response = await llm.ainvoke(user_content)
        ai_response = response.content

        # Try to extract JSON from the response
        json_match = re.search(r'\{.*\}', ai_response, re.DOTALL)
        if json_match:
            try:
                parsed_response = json.loads(json_match.group())
                
                # Apply AI Learning: Smart alternatives filtering
                item_analysis = analyze_item_types(ai_response)
                photo_count = len(files)
                smart_alternatives = get_smart_alternatives(item_analysis, lorry_data, photo_count)
                
                # Special handling for less than 2 photos with small items
                if photo_count < 2 and item_analysis.get("has_only_small_items", False):
                    # Force VAN as primary recommendation for very small moves
                    parsed_response["suggested_lorry_id"] = 4  # Van
                    # Set alternatives to 1-ton lorry only
                    parsed_response["alternative_lorries"] = [1]  # Small (1-Ton)
                    parsed_response["ai_learning"] = {
                        "applied": True,
                        "item_analysis": item_analysis,
                        "confidence_level": "very high",
                        "filtering_applied": "Van (1.5-Ton) primary + Small (1-Ton) alternatives",
                        "reasoning": f"Less than 2 photos detected with {item_analysis['small_item_count']} small item indicators. Automatically optimized for Van (1.5-Ton) as primary choice with Small (1-Ton) alternatives for maximum cost efficiency.",
                        "detection_method": "Photo count + enhanced context-aware analysis",
                        "photo_count_logic": f"Photo count: {photo_count} (triggered Van optimization)"
                    }
                # Regular smart filtering for 2+ photos
                elif item_analysis.get("has_only_small_items", False):
                    # Limit alternatives to 1-ton options for small items
                    if "alternative_lorries" in parsed_response:
                        # Filter alternatives to only include small/van options
                        original_alternatives = parsed_response["alternative_lorries"]
                        if isinstance(original_alternatives, list):
                            # Keep only alternatives that are in our smart alternatives list
                            filtered_alternatives = [alt for alt in original_alternatives if alt in smart_alternatives]
                            parsed_response["alternative_lorries"] = filtered_alternatives
                    
                    # Add detailed AI learning metadata
                    confidence_level = item_analysis.get("confidence", "low")
                    small_items = item_analysis.get("analysis_details", {}).get("small_items_detected", [])
                    
                    if confidence_level == "very high":
                        reasoning = f"Very high confidence: Detected {item_analysis['small_item_count']} small item indicators including {', '.join(small_items[:3])}. Strictly limited alternatives to 1-ton options for maximum cost optimization."
                    elif confidence_level == "high":
                        reasoning = f"High confidence: Detected {item_analysis['small_item_count']} small item indicators including {', '.join(small_items[:3])}. Limited alternatives to 1-ton and van options for cost optimization."
                    else:
                        reasoning = f"Medium confidence: Detected {item_analysis['small_item_count']} small item indicators. Limited alternatives to 1-ton options for cost optimization."
                    
                    parsed_response["ai_learning"] = {
                        "applied": True,
                        "item_analysis": item_analysis,
                        "confidence_level": confidence_level,
                        "filtering_applied": "1-ton alternatives only",
                        "reasoning": reasoning,
                        "detection_method": "Enhanced context-aware analysis",
                        "photo_count_logic": f"Photo count: {photo_count} (standard analysis)"
                    }
                else:
                    parsed_response["ai_learning"] = {
                        "applied": False,
                        "item_analysis": item_analysis,
                        "confidence_level": item_analysis.get("confidence", "low"),
                        "filtering_applied": "Full range alternatives",
                        "reasoning": "Mixed or large items detected. Full range of alternatives available for comprehensive choice.",
                        "detection_method": "Enhanced context-aware analysis"
                    }
                
                return parsed_response
            except json.JSONDecodeError:
                # If JSON parsing fails, return the raw response
                return {"raw_response": ai_response}
        else:
            return {"raw_response": ai_response}

    except Exception as e:
        raise HTTPException(500, detail=f"AI processing error: {str(e)}")

@router.post("/suggest-property")
async def suggest_property(property_data: Dict[str, Any] = Body(...)) -> Dict[str, Any]:
    """
    AI-powered property suggestion based on user preferences
    """
    if not llm:
        raise HTTPException(503, detail="AI service unavailable. Please configure GEMINI_API_KEY in .env.local file.")
    
    try:
        # Extract key property preferences
        budget = property_data.get('budget', 0)
        location = property_data.get('location', '')
        property_type = property_data.get('property_type', '')
        bedrooms = property_data.get('bedrooms', 0)
        amenities = property_data.get('amenities', [])
        
        # Create a prompt for property suggestions
        prompt = [
            {
                "role": "user",
                "content": f"""
                Based on these property preferences, suggest 3-5 properties:
                - Budget: RM {budget:,}
                - Location: {location}
                - Property Type: {property_type}
                - Bedrooms: {bedrooms}
                - Desired Amenities: {', '.join(amenities)}
                
                For each suggestion, provide:
                1. Property title
                2. Estimated price range
                3. Recommended location
                4. Key features
                5. Why it's a good match
                
                Return as JSON with: suggestions (array of property objects)
                """
            }
        ]
        
        response = await llm.ainvoke(prompt)
        ai_response = response.content
        
        # Try to extract JSON from the response
        json_match = re.search(r'\{.*\}', ai_response, re.DOTALL)
        if json_match:
            try:
                parsed_response = json.loads(json_match.group())
                return parsed_response
            except json.JSONDecodeError:
                return {"raw_response": ai_response}
        else:
            return {"raw_response": ai_response}
            
    except Exception as e:
        raise HTTPException(500, detail=f"Property suggestion error: {str(e)}")
