#!/usr/bin/env python3
"""
Test shot operations through the API to verify the new schema works correctly.
"""

import requests
import json
import sys

# API base URL
BASE_URL = "http://localhost:8000"

def get_auth_token():
    """Get authentication token for API testing."""
    # Try to login with admin credentials
    login_data = {
        "username": "admin",
        "password": "admin123"
    }
    
    try:
        response = requests.post(f"{BASE_URL}/auth/login", json=login_data)
        if response.status_code == 200:
            return response.json()["access_token"]
        else:
            print(f"❌ Login failed: {response.status_code} - {response.text}")
            return None
    except requests.exceptions.ConnectionError:
        print("❌ Cannot connect to API server. Please start the backend server first.")
        return None

def test_shot_api_operations():
    """Test shot operations through the API."""
    print("Testing shot API operations with new schema...")
    
    # Get authentication token
    token = get_auth_token()
    if not token:
        return False
    
    headers = {"Authorization": f"Bearer {token}"}
    
    try:
        # Test 1: Get shots and verify project_id is included
        print("\n1. Testing shot retrieval with project_id...")
        response = requests.get(f"{BASE_URL}/shots/", headers=headers)
        
        if response.status_code == 200:
            shots = response.json()
            if shots and len(shots) > 0:
                first_shot = shots[0]
                if "project_id" in first_shot:
                    print(f"✅ Shot API includes project_id: {first_shot['project_id']}")
                    print(f"   Shot: {first_shot['name']} (ID: {first_shot['id']})")
                else:
                    print("❌ Shot API response missing project_id field")
                    return False
            else:
                print("⚠️  No shots found in database")
        else:
            print(f"❌ Failed to retrieve shots: {response.status_code}")
            return False
        
        # Test 2: Test project filtering
        if shots and len(shots) > 0:
            test_project_id = shots[0]["project_id"]
            print(f"\n2. Testing project filtering (project_id={test_project_id})...")
            
            response = requests.get(f"{BASE_URL}/shots/?project_id={test_project_id}", headers=headers)
            if response.status_code == 200:
                filtered_shots = response.json()
                all_same_project = all(shot["project_id"] == test_project_id for shot in filtered_shots)
                if all_same_project:
                    print(f"✅ Project filtering works: {len(filtered_shots)} shots in project {test_project_id}")
                else:
                    print("❌ Project filtering failed: shots from different projects returned")
                    return False
            else:
                print(f"❌ Project filtering failed: {response.status_code}")
                return False
        
        # Test 3: Get episodes to test shot creation
        print("\n3. Testing shot creation with project_id validation...")
        response = requests.get(f"{BASE_URL}/episodes/", headers=headers)
        
        if response.status_code == 200:
            episodes = response.json()
            if episodes and len(episodes) > 0:
                test_episode = episodes[0]
                episode_id = test_episode["id"]
                episode_project_id = test_episode["project_id"]
                
                # Test creating a shot with correct project_id
                new_shot_data = {
                    "name": f"test_shot_migration_{episode_id}",
                    "episode_id": episode_id,
                    "description": "Test shot for migration verification",
                    "frame_start": 1001,
                    "frame_end": 1100
                }
                
                response = requests.post(f"{BASE_URL}/shots/", 
                                       json=new_shot_data, 
                                       headers=headers)
                
                if response.status_code == 200:
                    created_shot = response.json()
                    if created_shot["project_id"] == episode_project_id:
                        print(f"✅ Shot creation works: project_id auto-populated as {created_shot['project_id']}")
                        
                        # Clean up - delete the test shot
                        shot_id = created_shot["id"]
                        delete_response = requests.delete(f"{BASE_URL}/shots/{shot_id}", headers=headers)
                        if delete_response.status_code == 200:
                            print("✅ Test shot cleaned up successfully")
                        else:
                            print(f"⚠️  Failed to clean up test shot: {delete_response.status_code}")
                    else:
                        print(f"❌ Shot creation failed: wrong project_id {created_shot['project_id']} != {episode_project_id}")
                        return False
                else:
                    print(f"❌ Shot creation failed: {response.status_code} - {response.text}")
                    return False
            else:
                print("⚠️  No episodes found for testing shot creation")
        else:
            print(f"❌ Failed to retrieve episodes: {response.status_code}")
            return False
        
        # Test 4: Test shot update operations
        if shots and len(shots) > 0:
            print("\n4. Testing shot update operations...")
            test_shot = shots[0]
            shot_id = test_shot["id"]
            original_name = test_shot["name"]
            
            # Update shot name
            update_data = {
                "name": f"{original_name}_updated",
                "description": test_shot.get("description", ""),
                "frame_start": test_shot.get("frame_start", 1001),
                "frame_end": test_shot.get("frame_end", 1100)
            }
            
            response = requests.put(f"{BASE_URL}/shots/{shot_id}", 
                                  json=update_data, 
                                  headers=headers)
            
            if response.status_code == 200:
                updated_shot = response.json()
                if updated_shot["project_id"] == test_shot["project_id"]:
                    print(f"✅ Shot update works: project_id preserved as {updated_shot['project_id']}")
                    
                    # Restore original name
                    restore_data = update_data.copy()
                    restore_data["name"] = original_name
                    restore_response = requests.put(f"{BASE_URL}/shots/{shot_id}", 
                                                  json=restore_data, 
                                                  headers=headers)
                    if restore_response.status_code == 200:
                        print("✅ Shot name restored successfully")
                    else:
                        print(f"⚠️  Failed to restore shot name: {restore_response.status_code}")
                else:
                    print(f"❌ Shot update failed: project_id changed unexpectedly")
                    return False
            else:
                print(f"❌ Shot update failed: {response.status_code} - {response.text}")
                return False
        
        return True
        
    except requests.exceptions.ConnectionError:
        print("❌ Cannot connect to API server. Please start the backend server first.")
        return False
    except Exception as e:
        print(f"❌ API test failed with error: {e}")
        return False

if __name__ == "__main__":
    success = test_shot_api_operations()
    if success:
        print("\n✅ All shot API operations tests passed!")
        print("✅ Migration verification complete - all shot operations work correctly with new schema")
        sys.exit(0)
    else:
        print("\n❌ Some API tests failed!")
        sys.exit(1)