from agentlego.tools import BaseTool
import os
import cv2
import numpy as np
from insightface.app import FaceAnalysis

class FaceSimilarityTool(BaseTool):
    default_desc = 'Uses InsightFace to evaluate face similarity between two images.'
    
    def __init__(self, threshold=0.01, use_gpu=False):
        super().__init__()
        # Initialize InsightFace
        self.threshold = threshold
        
        # Use GPU if available (ctx_id=0) or CPU (ctx_id=-1)
        ctx_id = 0 if use_gpu else -1
        
        # Initialize FaceAnalysis
        print("Initializing InsightFace model...")
        self.app = FaceAnalysis(name='buffalo_s')
        self.app.prepare(ctx_id=ctx_id, det_size=(640, 640))
        print("✅ InsightFace model initialized successfully")
    
    def apply(self, img1_path: str, img2_path: str) -> str:
        try:
            # Validate file existence
            if not os.path.exists(img1_path):
                return f"Error: Image 1 not found at path: {img1_path}"
            if not os.path.exists(img2_path):
                return f"Error: Image 2 not found at path: {img2_path}"
            
            # Load images
            img1 = cv2.imread(img1_path)
            img2 = cv2.imread(img2_path)
            
            if img1 is None or img2 is None:
                return f"Error: Failed to load one or both images"
            
            # Get faces and embeddings
            faces1 = self.app.get(img1)
            faces2 = self.app.get(img2)
            
            # Check for no faces detected
            if len(faces1) == 0:
                return "Face similarity result: **No match** (No faces detected in first image)"
            if len(faces2) == 0:
                return "Face similarity result: **No match** (No faces detected in second image)"
            
            # Get first face from each image (use primary face if multiple detected)
            embedding1 = faces1[0].embedding
            embedding2 = faces2[0].embedding
            
            # Compute cosine similarity
            similarity = np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2))
            
            # Interpret result
            if similarity > self.threshold:
                result = "Match detected"
                match_text = "True"
            else:
                result = "No match"
                match_text = "False"
            
            # Format in a way similar to your original implementation
            return f"Face similarity result: **{match_text}** ({result}, similarity score: {similarity:.4f})"
        
        except Exception as e:
            return f"Error during face similarity computation: {str(e)}"