from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import requests
import torch
import easyocr
import numpy as np


class TagIdentifier:
    def __init__(self):
        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
        
        # Initialize OCR reader
        self.ocr_reader = easyocr.Reader(['en'])
        
        # Brand keywords for identification
        self.brand_keywords = [
            'nike', 'adidas', 'puma', 'reebok', 'champion', 'hanes',
            'fruit of the loom', 'gildan', 'jerzees', 'russell',
            'american apparel', 'next level', 'bella canvas'
        ]
        
    def load_image(self, image_url):
        try:
            response = requests.get(image_url, stream=True, timeout=10)
            response.raise_for_status()
            image = Image.open(response.raw)
            return image
        except Exception as e:
            raise ValueError(f"Failed to load image: {str(e)}")
    
    def get_image_features(self, image):
        """Get CLIP embeddings for the image"""
        inputs = self.processor(images=image, return_tensors="pt")
        with torch.no_grad():
            features = self.model.get_image_features(**inputs)
        return features.cpu().numpy().flatten()
    
    def extract_text_from_image(self, image):
        """Extract text from tag image using OCR"""
        try:
            # Convert PIL image to numpy array
            image_np = np.array(image)
            
            # Use EasyOCR to extract text
            results = self.ocr_reader.readtext(image_np)
            
            # Combine all detected text
            extracted_text = " ".join([result[1] for result in results])
            return extracted_text.lower()
            
        except Exception as e:
            print(f"OCR extraction failed: {e}")
            return ""
    
    def identify_brand(self, image):
        """Try to identify brand from image using both OCR and CLIP"""
        # OCR-based brand detection
        extracted_text = self.extract_text_from_image(image)
        
        for brand in self.brand_keywords:
            if brand in extracted_text:
                return brand, extracted_text
        
        # CLIP-based brand classification (if OCR fails)
        # Create text prompts for each brand
        brand_prompts = [f"a {brand} clothing tag" for brand in self.brand_keywords]
        brand_prompts.append("a generic clothing tag")
        
        # Get image and text features
        inputs = self.processor(
            text=brand_prompts, 
            images=image, 
            return_tensors="pt", 
            padding=True
        )
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = outputs.logits_per_image.softmax(dim=1)
            
        # Get best match
        best_idx = probs.argmax().item()
        if best_idx < len(self.brand_keywords):
            return self.brand_keywords[best_idx], extracted_text
        
        return "unknown", extracted_text
