import os
import random
import requests
import time
import psycopg2
from pathlib import Path

# Configuration
BASE_URL = "http://localhost:5000"
DB_NAME = "morphguard"
DB_USER = "morphguard"
DB_PASS = "morphguard"
DB_HOST = "localhost"

# Update credentials if changed in your env
ADMIN_USER = "admin"
ADMIN_PASS = "morphguard_admin" 

DATA_DIRS = {
    "real": "/home/juanquy/dev/MorphGuard/data/datasets/morph_detection/real",
    "morph": "/home/juanquy/dev/MorphGuard/data/datasets/morph_detection/morph",
    "pexels": "/home/juanquy/dev/MorphGuard/data/real_faces/pexels"
}

def get_db_connection():
    try:
        conn = psycopg2.connect(
            dbname=DB_NAME,
            user=DB_USER,
            password=DB_PASS,
            host=DB_HOST
        )
        return conn
    except Exception as e:
        print(f"Error connecting to DB: {e}")
        return None

def get_initial_metric_count(conn):
    if not conn: return 0
    try:
        cur = conn.cursor()
        cur.execute("SELECT COUNT(*) FROM detection_metrics;")
        count = cur.fetchone()[0]
        cur.close()
        return count
    except Exception as e:
        print(f"Error querying metrics: {e}")
        return 0

def collect_images(num_per_category=8):
    images = []
    print("Collecting images...")
    for category, path in DATA_DIRS.items():
        if not os.path.exists(path):
            print(f"Warning: Path not found: {path}")
            continue
        
        # Recursive search for images
        all_files = []
        for root, dirs, files in os.walk(path):
            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    all_files.append(os.path.join(root, file))
        
        if not all_files:
            print(f"Warning: No images found in {path}")
            continue
            
        selected = random.sample(all_files, min(len(all_files), num_per_category))
        print(f"  - {category}: {len(selected)} images selected")
        images.extend([(category, img) for img in selected])
    
    random.shuffle(images)
    return images

def test_api(images):
    print(f"\nStarting API testing with {len(images)} images...")
    success_count = 0
    session = requests.Session()
    
    # Login first
    print("Attempting login...")
    try:
        login_resp = session.post(f"{BASE_URL}/login", data={'username': ADMIN_USER, 'password': ADMIN_PASS})
        if login_resp.status_code != 200 and login_resp.url.endswith('/login'): # Redirect back to login means fail
             # Try JSON login just in case
            login_resp = session.post(f"{BASE_URL}/login", json={'username': ADMIN_USER, 'password': ADMIN_PASS})
            
        if login_resp.status_code != 200 and 'Invalid credentials' in login_resp.text:
             print("LOGIN FAILED: Check admin credentials.")
             return 0
        print("Login successful.")
    except Exception as e:
        print(f"Login error: {e}")
        return 0
    
    for i, (category, img_path) in enumerate(images, 1):
        print(f"[{i}/{len(images)}] Testing {category} image: {os.path.basename(img_path)}", end=" ... ")
        
        try:
            with open(img_path, 'rb') as f:
                start_time = time.time()
                # Use session which stores cookie
                response = session.post(
                    f"{BASE_URL}/api/detect",
                    files={'image': f}, # ERROR FIX: detection.py expects 'image', not 'file'
                    data={'mode': 'deep_forensics'} # Trigger forensics
                )
                duration = time.time() - start_time
                
            if response.status_code == 200:
                result = response.json()
                # Check for standard response keys
                if 'is_morphed' in result:
                    print(f"SUCCESS ({duration:.2f}s) - Morphed: {result['is_morphed']}, Conf: {result.get('confidence', 0):.2f}")
                    success_count += 1
                else:
                    print(f"FAILED - Invalid JSON response: {result.keys()}")
            else:
                print(f"FAILED - Status {response.status_code}: {response.text[:100]}")
                
        except Exception as e:
            print(f"ERROR: {e}")
            
    return success_count

def verify_db_metrics(initial_count):
    print("\nVerifying Database Metrics...")
    conn = get_db_connection()
    if not conn:
        print("Skipping DB verification due to connection error")
        return
        
    final_count = get_initial_metric_count(conn)
    conn.close()
    
    diff = final_count - initial_count
    print(f"Metrics count: Initial={initial_count}, Final={final_count}, Added={diff}")
    
    if diff > 0:
        print("✅ DATABASE VERIFICATION PASSED: Metrics were recorded.")
    else:
        print("❌ DATABASE VERIFICATION FAILED: No new metrics found.")

def main():
    # Wait for server to start if needed
    print("Checking server status...")
    try:
        requests.get(BASE_URL)
    except:
        print("Waiting for server to start...")
        time.sleep(5)
    
    conn = get_db_connection()
    initial_metrics = get_initial_metric_count(conn)
    if conn: conn.close()
    
    images = collect_images(8) # 8 * 3 = 24 images
    
    if not images:
        print("No images found to test! Check paths.")
        return
        
    successful_requests = test_api(images)
    
    print(f"\nTest Complete: {successful_requests}/{len(images)} requests successful.")
    
    verify_db_metrics(initial_metrics)

if __name__ == "__main__":
    main()