import requests
import json
import os

# Get API credentials from environment variables
EPA_AQS_API_BASE_URL = "https://aqs.epa.gov/data/api"
EMAIL = os.environ.get("EPA_AQS_EMAIL", "")
API_KEY = os.environ.get("EPA_AQS_API_KEY", "")

def print_separator():
    print("=" * 50)

def test_api_connection():
    """Test the basic API connectivity and credentials"""
    print_separator()
    print("TESTING API CONNECTION")
    print(f"Using email: {EMAIL}")
    print(f"API key provided: {'Yes' if API_KEY else 'No'}")
    
    if not EMAIL or not API_KEY:
        print("WARNING: Email or API key is missing. Real API calls will fail.")
        return False
    
    # Test with a simple API call
    endpoint = f"{EPA_AQS_API_BASE_URL}/list/states"
    params = {
        "email": EMAIL,
        "key": API_KEY
    }
    
    try:
        response = requests.get(endpoint, params=params)
        print(f"Response status code: {response.status_code}")
        data = response.json()
        
        # Print the entire response structure for examination
        print("First 500 characters of the response structure:")
        response_str = json.dumps(data, indent=2)
        print(response_str[:500] + "..." if len(response_str) > 500 else response_str)
        
        # Specific handling for the observed response structure
        if isinstance(data, dict) and "Header" in data and isinstance(data["Header"], list) and len(data["Header"]) > 0:
            header = data["Header"][0]
            if header.get("status") == "Success":
                print("API connection successful!")
                return True
        
        print("API connection failed: Response format not as expected")
        print("Full response keys:", list(data.keys()) if isinstance(data, dict) else "Not a dictionary")
        return False
    except Exception as e:
        print(f"Exception during API connection test: {e}")
        return False

def test_state_lookup(state_code):
    """Test looking up a specific state"""
    print_separator()
    print(f"TESTING STATE LOOKUP FOR: {state_code}")
    
    endpoint = f"{EPA_AQS_API_BASE_URL}/list/states"
    params = {
        "email": EMAIL,
        "key": API_KEY
    }
    
    try:
        response = requests.get(endpoint, params=params)
        data = response.json()
        
        # Handle the specific response structure we observed
        if isinstance(data, dict) and "Data" in data and isinstance(data["Data"], list):
            states = data["Data"]
            print(f"Found {len(states)} states in the API")
            
            # Look for the specific state
            matching_states = [s for s in states if str(s.get("code")) == str(state_code) or s.get("value_represented") == state_code]
            if matching_states:
                print(f"Found matching state: {matching_states[0]}")
                return matching_states[0]
            else:
                print(f"No matching state found for code: {state_code}")
                print("Available state codes (first 10):")
                for s in states[:10]:  # Print first 10 for brevity
                    print(f"  {s.get('code')} - {s.get('value_represented')}")
                if len(states) > 10:
                    print(f"  ... and {len(states)-10} more")
                return None
        else:
            print("API lookup failed: Unexpected response format")
            return None
    except Exception as e:
        print(f"Exception during state lookup: {e}")
        return None

def test_counties_lookup(state_code):
    """Test looking up counties for a state"""
    print_separator()
    print(f"TESTING COUNTIES LOOKUP FOR STATE: {state_code}")
    
    endpoint = f"{EPA_AQS_API_BASE_URL}/list/countiesByState"
    params = {
        "email": EMAIL,
        "key": API_KEY,
        "state": state_code
    }
    
    try:
        response = requests.get(endpoint, params=params)
        data = response.json()
        
        # Print the beginning of the response structure for examination
        print("First 500 characters of the response structure:")
        response_str = json.dumps(data, indent=2)
        print(response_str[:500] + "..." if len(response_str) > 500 else response_str)
        
        counties = []
        if isinstance(data, dict) and "Data" in data and isinstance(data["Data"], list):
            counties = data["Data"]
        
        print(f"Found {len(counties)} counties for state {state_code}")
        
        if counties:
            print("Sample counties (first 5):")
            for c in counties[:5]:  # Print first 5 for brevity
                print(f"  {c}")
            if len(counties) > 5:
                print(f"  ... and {len(counties)-5} more")
        else:
            print("No counties found for this state")
        
        return counties
    except Exception as e:
        print(f"Exception during counties lookup: {e}")
        return None

def test_monitors_lookup(state_code):
    """Test looking up monitors for a state"""
    print_separator()
    print(f"TESTING MONITORS LOOKUP FOR STATE: {state_code}")
    
    endpoint = f"{EPA_AQS_API_BASE_URL}/monitors/byState"
    params = {
        "email": EMAIL,
        "key": API_KEY,
        "state": state_code,
        "bdate": "20240101",
        "edate": "20240414"
    }
    
    try:
        response = requests.get(endpoint, params=params)
        data = response.json()
        
        # Print the beginning of the response structure for examination
        print("First 500 characters of the response structure:")
        response_str = json.dumps(data, indent=2)
        print(response_str[:500] + "..." if len(response_str) > 500 else response_str)
        
        monitors = []
        if isinstance(data, dict) and "Data" in data and isinstance(data["Data"], list):
            monitors = data["Data"]
        
        print(f"Found {len(monitors)} monitors for state {state_code}")
        
        if monitors and len(monitors) > 0:
            print("Sample monitor:")
            sample_monitor = monitors[0]
            print(json.dumps(sample_monitor, indent=2))
            
            # Check keys that might be causing issues
            print("\nChecking critical keys for indexing issues:")
            expected_keys = ["state_code", "county_code", "site_number", "parameter_code", "parameter_name", "latitude", "longitude", "local_site_name"]
            for key in expected_keys:
                if key in sample_monitor:
                    print(f"  ✓ '{key}' found: {sample_monitor[key]}")
                else:
                    print(f"  ✗ '{key}' missing!")
        else:
            print("No monitors found for this state")
        
        return monitors
    except Exception as e:
        print(f"Exception during monitors lookup: {e}")
        return None

def run_tests():
    """Run all tests"""
    print("Starting API Tests")
    api_working = test_api_connection()
    
    if api_working:
        # Test some states with different formats
        for state_code in ["06", "36", "CA", "NY"]:
            state_info = test_state_lookup(state_code)
            if state_info:
                # Use the state code according to what was found
                state_code_to_use = state_info.get("code")
                if state_code_to_use:
                    test_counties_lookup(state_code_to_use)
                    test_monitors_lookup(state_code_to_use)
                else:
                    print(f"Could not determine the proper state code format for {state_code}")
    
    print_separator()
    print("Tests completed.")

if __name__ == "__main__":
    run_tests()