import os
import sys
import torch
import time
from PIL import Image

# Add current workspace to path
sys.path.append(os.getcwd())

# Ensure output directory exists
os.makedirs('scratch', exist_ok=True)

def main():
    print("====================================================")
    print("Starting MorphGuard Model & Integration Diagnostics")
    print("====================================================")
    
    # 1. Initialize API
    print("\n[Step 1] Initializing MorphGuardAPI...")
    from morphguard_api import MorphGuardAPI
    
    # Set CTM to enabled in environment
    os.environ['CTM_ENABLED'] = 'true'
    
    api = MorphGuardAPI(
        detector_path='models/morph_detector.pth',
        demorph_path='models/demorpher.pth'
    )
    print(f"API initialized on device: {api.device}")
    
    # Check loaded status
    print(f"Detector path: {api.config['detector_path']} (Exists: {os.path.exists(api.config['detector_path'])})")
    print(f"Demorpher path: {api.config['demorph_path']} (Exists: {os.path.exists(api.config['demorph_path'])})")
    print(f"CTM path: models/ctm_forensic.pth (Exists: {os.path.exists('models/ctm_forensic.pth')})")
    
    # 2. Test Fast Detection on Real Image
    real_img_path = 'data/datasets/morph_detection/real/001_03.jpg'
    print(f"\n[Step 2] Testing Fast Detection (Tier 1) on Real Image: {real_img_path}")
    if os.path.exists(real_img_path):
        start = time.time()
        res_real = api.detect_morph(real_img_path)
        elapsed = (time.time() - start) * 1000
        print(f"Result -> Is Morphed: {res_real.get('is_morphed')}")
        print(f"Confidence: {res_real.get('confidence'):.4f}")
        print(f"Time Taken: {elapsed:.2f}ms")
    else:
        print(f"Error: Real image {real_img_path} not found!")

    # 3. Test Fast Detection on Morphed Image
    morph_img_path = 'data/datasets/morph_detection/morph/dim_a/001_010.png'
    print(f"\n[Step 3] Testing Fast Detection (Tier 1) on Morphed Image: {morph_img_path}")
    if os.path.exists(morph_img_path):
        start = time.time()
        res_morph = api.detect_morph(morph_img_path)
        elapsed = (time.time() - start) * 1000
        print(f"Result -> Is Morphed: {res_morph.get('is_morphed')}")
        print(f"Confidence: {res_morph.get('confidence'):.4f}")
        print(f"Time Taken: {elapsed:.2f}ms")
    else:
        print(f"Error: Morphed image {morph_img_path} not found!")

    # 4. Test CTM Forensic Analysis (Tier 2) on Morphed Image
    print(f"\n[Step 4] Testing Deep Forensic Analysis (Tier 2 - CTM) on Morphed Image: {morph_img_path}")
    if os.path.exists(morph_img_path):
        from src.models.ctm_forensic_agent import get_ctm_agent
        ctm_agent = get_ctm_agent()
        
        start = time.time()
        ctm_result = ctm_agent.analyze(
            morph_img_path,
            generate_evidence=True,
            request_id='diag_test'
        )
        elapsed = (time.time() - start) * 1000
        
        print(f"Result -> Is Morphed: {ctm_result.is_morphed}")
        print(f"Confidence: {ctm_result.confidence:.4f}")
        print(f"Forensic Steps: {ctm_result.forensic_steps}")
        print(f"Attention Regions: {ctm_result.attention_regions}")
        print(f"Evidence Video GIF Path: {ctm_result.evidence_video_path}")
        print(f"Data Tier: {ctm_result.tier}")
        print(f"Time Taken: {elapsed:.2f}ms")
    else:
        print(f"Error: Morphed image {morph_img_path} not found!")

    # 5. Test Face Demorphing
    output_demorph_path = 'scratch/demorphed_001_010.png'
    print(f"\n[Step 5] Testing Transformer Demorpher on Morphed Image: {morph_img_path}")
    if os.path.exists(morph_img_path):
        ref_path = 'data/datasets/morph_detection/real/001_03.jpg'
        print(f"Using reference face: {ref_path}")
        
        start = time.time()
        demorph_res = api.demorph_image(
            image_path=morph_img_path,
            reference_path=ref_path if os.path.exists(ref_path) else None,
            output_path=output_demorph_path,
            method='transformer'
        )
        elapsed = (time.time() - start) * 1000
        
        print(f"Demorph Result -> Success: {demorph_res.get('success')}")
        print(f"Output Saved To: {demorph_res.get('output_path')}")
        print(f"Processing Time: {elapsed:.2f}ms")
        
        if os.path.exists(output_demorph_path):
            img_size = os.path.getsize(output_demorph_path)
            print(f"Demorphed File Verified: Size = {img_size} bytes")
            # Open image to verify it's a valid PIL image
            try:
                with Image.open(output_demorph_path) as img:
                    print(f"Demorphed Image format: {img.format}, Mode: {img.mode}, Size: {img.size}")
                print("✅ Face Demorpher works correctly!")
            except Exception as e:
                print(f"❌ Error verifying output image: {e}")
        else:
            print("❌ Error: Demorphed image file was not created!")
    else:
        print(f"Error: Morphed image {morph_img_path} not found!")

    print("\n====================================================")
    print("Diagnostics Finished Successfully!")
    print("====================================================")

if __name__ == '__main__':
    main()