import torch
import numpy as np
from sentence_transformers import SentenceTransformer
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import random
import logging
import json
import os

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Chatbot:
    def __init__(self):
        self.device = 'cpu'  # Force CPU usage for Hugging Face Spaces
        logger.info("🚀 Initializing Fashion Chatbot with CPU...")
        self.model = None
        self.product_data = {}
        self.images = {}
        self.product_embeddings = None
        self.load_models()
        self.setup_sample_data()
    
    def load_models(self):
        """Load all required models with CPU-only configuration"""
        try:
            logger.info("📥 Loading SentenceTransformer model on CPU...")
            
            # Force CPU for all operations
            torch.device('cpu')
            
            # Load a lightweight model suitable for CPU
            self.model = SentenceTransformer(
                'all-MiniLM-L6-v2',  # Lightweight model for CPU
                device='cpu'
            )
            
            logger.info("✅ Model loaded successfully on CPU")
            
        except Exception as e:
            logger.error(f"❌ Error loading model: {e}")
            # Create a dummy model for fallback
            self.model = None
    
    def setup_sample_data(self):
        """Setup sample fashion product data for demonstration"""
        logger.info("🛍️ Setting up sample fashion data...")
        
        # Sample fashion products data
        self.product_data = {
            0: {
                'productDisplayName': 'Classic White T-Shirt',
                'masterCategory': 'Apparel',
                'articleType': 'T-Shirt',
                'usage': 'Casual',
                'season': 'All Season',
                'gender': 'Unisex',
                'baseColour': 'White',
                'price': 29.99
            },
            1: {
                'productDisplayName': 'Denim Jacket',
                'masterCategory': 'Apparel',
                'articleType': 'Jacket',
                'usage': 'Casual',
                'season': 'Spring, Fall',
                'gender': 'Unisex',
                'baseColour': 'Blue',
                'price': 89.99
            },
            2: {
                'productDisplayName': 'Black Leather Boots',
                'masterCategory': 'Footwear',
                'articleType': 'Boots',
                'usage': 'Casual',
                'season': 'Winter, Fall',
                'gender': 'Unisex',
                'baseColour': 'Black',
                'price': 129.99
            },
            3: {
                'productDisplayName': 'Summer Floral Dress',
                'masterCategory': 'Apparel',
                'articleType': 'Dress',
                'usage': 'Casual',
                'season': 'Summer',
                'gender': 'Women',
                'baseColour': 'Multicolor',
                'price': 59.99
            },
            4: {
                'productDisplayName': 'Sports Running Shoes',
                'masterCategory': 'Footwear',
                'articleType': 'Sports Shoes',
                'usage': 'Sports',
                'season': 'All Season',
                'gender': 'Unisex',
                'baseColour': 'White',
                'price': 79.99
            },
            5: {
                'productDisplayName': 'Wool Winter Scarf',
                'masterCategory': 'Accessories',
                'articleType': 'Scarf',
                'usage': 'Casual',
                'season': 'Winter',
                'gender': 'Unisex',
                'baseColour': 'Grey',
                'price': 34.99
            }
        }
        
        # Generate sample product images
        self.images = {}
        for pid in self.product_data.keys():
            self.images[pid] = self.generate_sample_image(pid)
        
        # Create sample embeddings for products
        self.create_sample_embeddings()
        
        logger.info(f"✅ Loaded {len(self.product_data)} sample products")
    
    def generate_sample_image(self, product_id):
        """Generate a sample product image for demonstration"""
        # Create a simple colored image with text
        img = Image.new('RGB', (200, 200), color=self.get_color_for_product(product_id))
        draw = ImageDraw.Draw(img)
        
        # Add product type text
        product_type = self.product_data[product_id]['articleType']
        draw.text((50, 90), product_type, fill='white')
        
        return img
    
    def get_color_for_product(self, product_id):
        """Get color based on product"""
        color_map = {
            'White': (255, 255, 255),
            'Blue': (0, 0, 255),
            'Black': (0, 0, 0),
            'Multicolor': (255, 0, 0),
            'Grey': (128, 128, 128)
        }
        base_color = self.product_data[product_id]['baseColour']
        return color_map.get(base_color, (200, 200, 200))
    
    def create_sample_embeddings(self):
        """Create sample embeddings for products"""
        try:
            if self.model is not None:
                product_descriptions = []
                for pid, data in self.product_data.items():
                    desc = f"{data['productDisplayName']} {data['articleType']} {data['usage']} {data['season']} {data['gender']}"
                    product_descriptions.append(desc)
                
                self.product_embeddings = self.model.encode(product_descriptions)
            else:
                # Create dummy embeddings
                self.product_embeddings = np.random.randn(len(self.product_data), 384)
        except Exception as e:
            logger.error(f"Error creating embeddings: {e}")
            self.product_embeddings = np.random.randn(len(self.product_data), 384)
    
    def load_data(self):
        """Load product data - using sample data for demo"""
        logger.info("📊 Loading product data...")
        # Data is already loaded in setup_sample_data
        pass
    
    def generate_image_caption(self, image_path):
        """Generate caption for uploaded image"""
        try:
            # For CPU deployment, use a simpler approach
            image = Image.open(image_path)
            
            # Simple analysis based on image characteristics
            width, height = image.size
            dominant_color = self.get_dominant_color(image)
            
            # Generate descriptive caption based on image properties
            size_desc = "large" if width > 1000 else "medium" if width > 500 else "small"
            color_desc = self.get_color_name(dominant_color)
            
            captions = [
                f"A {size_desc} {color_desc} fashion item perfect for your style",
                f"Stylish {color_desc} clothing item that matches current trends",
                f"Fashionable {size_desc} apparel in {color_desc} color",
                f"Trendy {color_desc} fashion piece suitable for various occasions"
            ]
            
            return random.choice(captions)
            
        except Exception as e:
            logger.error(f"Error generating caption: {e}")
            return "A fashionable clothing item that suits your style"
    
    def get_dominant_color(self, image):
        """Get dominant color from image (simplified)"""
        try:
            # Resize image for faster processing
            image = image.resize((50, 50))
            # Convert to numpy array and get average color
            np_image = np.array(image)
            return tuple(np.mean(np_image, axis=(0, 1)).astype(int))
        except:
            return (128, 128, 128)  # Default gray
    
    def get_color_name(self, rgb):
        """Convert RGB to color name"""
        colors = {
            (255, 255, 255): "white",
            (0, 0, 0): "black",
            (255, 0, 0): "red",
            (0, 255, 0): "green",
            (0, 0, 255): "blue",
            (255, 255, 0): "yellow",
            (128, 128, 128): "gray",
            (255, 165, 0): "orange",
            (128, 0, 128): "purple"
        }
        
        # Find closest color
        min_dist = float('inf')
        closest_color = "colored"
        for color, name in colors.items():
            dist = sum((a - b) ** 2 for a, b in zip(rgb, color))
            if dist < min_dist:
                min_dist = dist
                closest_color = name
                
        return closest_color
    
    def generate_response(self, query):
        """Generate chatbot response and recommendations"""
        try:
            # Fashion-related responses
            fashion_responses = {
                'casual': "Great choice! Casual wear is perfect for everyday comfort and style.",
                'formal': "Elegant choice! Formal wear always makes a strong impression.",
                'sports': "Active lifestyle! Sports wear combines comfort and performance.",
                'summer': "Perfect for warm weather! Light and breathable fabrics work best.",
                'winter': "Stay warm and stylish! Layering is key for winter fashion.",
                'dress': "Dresses are versatile and always in style!",
                'shirt': "Classic shirts never go out of fashion!",
                'shoes': "The right shoes can complete any outfit!",
                'jacket': "Jackets add style and functionality to any outfit!"
            }
            
            # Generate contextual response
            query_lower = query.lower()
            response_key = None
            
            for key in fashion_responses.keys():
                if key in query_lower:
                    response_key = key
                    break
            
            if response_key:
                bot_response = fashion_responses[response_key]
            else:
                generic_responses = [
                    f"I found some great fashion items related to '{query}'!",
                    f"Based on your interest in '{query}', here are my recommendations:",
                    f"Here are some stylish options for '{query}':",
                    f"Perfect! I have some fashion suggestions for '{query}':"
                ]
                bot_response = random.choice(generic_responses)
            
            # Get recommendations
            recommended_products = self.get_recommendations(query)
            
            return bot_response, recommended_products
            
        except Exception as e:
            logger.error(f"Error generating response: {e}")
            return "I apologize, but I'm having trouble processing your request right now.", []
    
    def get_recommendations(self, query, top_k=3):
        """Get product recommendations based on query"""
        try:
            if self.model is not None and self.product_embeddings is not None:
                # Encode query
                query_embedding = self.model.encode([query])
                
                # Calculate similarities (using dot product for simplicity)
                similarities = np.dot(self.product_embeddings, query_embedding.T).flatten()
                
                # Get top products
                top_indices = np.argsort(similarities)[::-1][:top_k]
            else:
                # Fallback: random recommendations
                top_indices = random.sample(list(self.product_data.keys()), min(top_k, len(self.product_data)))
            
            recommended_products = []
            for idx in top_indices:
                recommended_products.append({
                    'corpus_id': idx,
                    'score': 0.9 - (len(recommended_products) * 0.1)
                })
            
            return recommended_products
            
        except Exception as e:
            logger.error(f"Error getting recommendations: {e}")
            # Return random products as fallback
            return [{'corpus_id': i, 'score': 0.8} for i in range(min(3, len(self.product_data)))]
    
    def get_product_info(self, product_id):
        """Get complete product information"""
        try:
            if product_id in self.product_data:
                data = self.product_data[product_id]
                return {
                    'name': data['productDisplayName'],
                    'category': data['masterCategory'],
                    'article_type': data['articleType'],
                    'usage': data['usage'],
                    'season': data['season'],
                    'gender': data['gender'],
                    'color': data['baseColour'],
                    'price': data['price'],
                    'image': self.images.get(product_id)
                }
            return None
        except Exception as e:
            logger.error(f"Error getting product info: {e}")
            return None