#!/usr/bin/env python3
"""
Simple Face Identifier for PunimTag
Works with the punimtag_simple.db database
"""

import sqlite3
import os
from PIL import Image
import pickle

class SimpleFaceIdentifier:
    def __init__(self, db_path='punimtag_simple.db'):
        self.db_path = db_path
        
    def get_unidentified_faces(self, limit=10):
        """Get a limited number of unidentified faces"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        
        c.execute('''SELECT f.id, f.image_id, i.path, i.filename, f.top, f.right, f.bottom, f.left
                     FROM faces f
                     JOIN images i ON f.image_id = i.id
                     WHERE f.person_id IS NULL
                     LIMIT ?''', (limit,))
        
        faces = c.fetchall()
        conn.close()
        return faces
    
    def add_person(self, name):
        """Add a new person"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute('INSERT OR IGNORE INTO people (name) VALUES (?)', (name,))
        c.execute('SELECT id FROM people WHERE name = ?', (name,))
        person_id = c.fetchone()[0]
        conn.commit()
        conn.close()
        return person_id
    
    def assign_face(self, face_id, person_id):
        """Assign a face to a person"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute('UPDATE faces SET person_id = ?, is_confirmed = 1 WHERE id = ?', 
                  (person_id, face_id))
        conn.commit()
        conn.close()
    
    def run_cli_identifier(self):
        """Run command line identifier"""
        print("\n🏷️  Simple Face Identifier")
        print("=" * 50)
        
        faces = self.get_unidentified_faces(50)  # Get first 50 faces
        
        if not faces:
            print("No unidentified faces found!")
            return
        
        print(f"Found {len(faces)} unidentified faces to process...")
        print("For each face, enter the person's name or 's' to skip\n")
        
        for i, (face_id, image_id, path, filename, top, right, bottom, left) in enumerate(faces):
            print(f"\nFace {i+1}/{len(faces)}")
            print(f"📁 File: {filename}")
            print(f"📍 Location: top={top}, right={right}, bottom={bottom}, left={left}")
            
            # Try to display basic info about the image
            try:
                if os.path.exists(path):
                    with Image.open(path) as img:
                        print(f"🖼️  Image size: {img.size}")
                else:
                    print("⚠️  Image file not found")
            except Exception as e:
                print(f"⚠️  Could not read image: {e}")
            
            while True:
                name = input(f"👤 Who is this person? (or 's' to skip): ").strip()
                
                if name.lower() == 's':
                    print("⏭️  Skipped")
                    break
                elif name:
                    try:
                        person_id = self.add_person(name)
                        self.assign_face(face_id, person_id)
                        print(f"✅ Identified as '{name}'")
                        break
                    except Exception as e:
                        print(f"❌ Error: {e}")
                else:
                    print("Please enter a name or 's' to skip")
        
        print(f"\n🎉 Completed processing {len(faces)} faces!")
        
        # Show remaining count
        remaining = self.get_remaining_count()
        if remaining > 0:
            print(f"📊 {remaining} unidentified faces remaining")
            print("Run the script again to continue identifying faces")
        else:
            print("🏆 All faces have been identified!")
    
    def get_remaining_count(self):
        """Get count of remaining unidentified faces"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute('SELECT COUNT(*) FROM faces WHERE person_id IS NULL')
        count = c.fetchone()[0]
        conn.close()
        return count

if __name__ == "__main__":
    identifier = SimpleFaceIdentifier()
    identifier.run_cli_identifier()