# stress_test.py
import requests
import time
import os
from dotenv import load_dotenv
import logging
from datetime import datetime
import threading # Import threading

# Configure logging for the script
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('stress_test.log')
    ]
)
logger = logging.getLogger(__name__)

# Load environment variables from .env file
load_dotenv()

# --- Configuration ---
BASE_URL = os.getenv("SERVER_BASE_URL", "http://localhost:8000") # Defaulting back to http for local dev unless overridden
CREATE_SONG_ENDPOINT = f"{BASE_URL}/create_song"
# Use Basic Auth credentials
USERNAME = "user"
PASSWORD = os.getenv("HTTP_BASIC_PASSWORD")
assert PASSWORD, "HTTP_BASIC_PASSWORD environment variable not set in .env file or environment. Exiting."
REQUEST_INTERVAL_SECONDS = 120
SONG_TOPIC = "Automated latency test song about a happy little cloud"
SONG_MODEL = "chirp-v4-h-api"
# Timeout for connecting and receiving headers (seconds)
REQUEST_TIMEOUT = 60
# Timeout for waiting for stream events (audio/image) (seconds)
STREAM_TIMEOUT = 180 # Increased timeout for the streaming part
# --- End Configuration ---

def make_request():
    """Makes a single POST request to the /create_song endpoint and waits for audio/image events."""
    payload = {
        "topic": SONG_TOPIC,
        "model": SONG_MODEL
    }
    auth = (USERNAME, PASSWORD)
    response = None # Initialize response to None
    start_time = time.time()

    try:
        logger.info(f"Sending request to {CREATE_SONG_ENDPOINT} with topic: '{SONG_TOPIC}'")
        # Use stream=True to get headers quickly and manage the connection manually
        response = requests.post(
            CREATE_SONG_ENDPOINT,
            json=payload,
            auth=auth,
            stream=True, # Get headers only first
            timeout=REQUEST_TIMEOUT # Apply timeout ONLY to connection and header retrieval
        )
        response.raise_for_status() # Check if the initial connection and request were successful (2xx status)
        logger.info(f"Request initiated successfully (Status Code: {response.status_code}). Backend processing started, waiting for stream events...")

        # Now, consume the stream and wait for specific events
        received_image = False
        received_audio = False
        stream_start_time = time.time()

        # iter_lines handles decoding and line splitting
        for line_bytes in response.iter_lines(decode_unicode=True):
            # Check for stream timeout
            if time.time() - stream_start_time > STREAM_TIMEOUT:
                 logger.error(f"Stream timed out after {STREAM_TIMEOUT} seconds while waiting for audio/image events.")
                 break # Exit the loop on timeout

            if line_bytes: # Filter out keep-alive newlines
                logger.debug(f"Received line: {line_bytes}")
                if line_bytes.strip() == "event: image_generated":
                    received_image = True
                    logger.info("Received 'image_generated' event.")
                elif line_bytes.strip() == "event: gen_streaming":
                    received_audio = True
                    logger.info("Received 'gen_streaming' event.")

                # Check if both events have been received
                if received_image and received_audio:
                    elapsed_time = time.time() - start_time
                    logger.info(f"Received both image and audio stream events after {elapsed_time:.2f} seconds. Test successful.")
                    break # Exit the loop successfully

        # After loop: Check if we exited due to success or timeout/stream end
        if not (received_image and received_audio):
             # If the loop finished without finding both events (e.g., timeout or stream ended prematurely)
             elapsed_time = time.time() - start_time
             logger.warning(f"Stream finished or timed out after {elapsed_time:.2f} seconds without receiving both required events (Image: {received_image}, Audio: {received_audio}).")


    except requests.exceptions.Timeout:
        logger.error(f"Request timed out after {REQUEST_TIMEOUT} seconds (connection/headers phase).")
    except requests.exceptions.RequestException as e:
        logger.error(f"Request failed: {e}")
        # Log response body if available on error
        if response is not None and response.text:
            logger.error(f"Response body: {response.text[:500]}...") # Log first 500 chars
    except Exception as e:
        logger.error(f"An unexpected error occurred during the request: {e}", exc_info=True)
    finally:
        # IMPORTANT: Close the response *after* processing the stream or encountering an error
        if response:
            response.close()
            logger.debug("Response closed.")


if __name__ == "__main__":
    # Ensure password is set
    if not PASSWORD:
        logger.error("HTTP_BASIC_PASSWORD environment variable not set in .env file or environment. Exiting.")
        exit(1)

    # Check if BASE_URL is still default and maybe warn if it's not HTTPS for production testing
    if BASE_URL == "http://localhost:8000":
        logger.warning("Using default BASE_URL 'http://localhost:8000'. Ensure this is correct for your test.")
    elif not BASE_URL.startswith("https://") and "localhost" not in BASE_URL:
         logger.warning(f"BASE_URL '{BASE_URL}' does not use HTTPS. Consider using HTTPS for production endpoints.")


    logger.info("Starting stress test script...")
    logger.info(f"Target URL: {CREATE_SONG_ENDPOINT}")
    logger.info(f"Request interval: {REQUEST_INTERVAL_SECONDS} seconds")
    logger.info(f"Request timeout: {REQUEST_TIMEOUT} seconds")

    while True:
        # Create and start a new thread for each request
        thread = threading.Thread(target=make_request, daemon=True) # Use daemon threads
        thread.start()
        logger.info(f"Started new request thread. Sleeping for {REQUEST_INTERVAL_SECONDS} seconds...")
        time.sleep(REQUEST_INTERVAL_SECONDS)