import runpod
import httpx
import json
import os
import subprocess
import sys
import time

# Configuration for the FastAPI server running locally within the container
FASTAPI_SERVER_URL = "http://127.0.0.1:8004/tts"

# Get Runpod AI API Key from environment variable if needed for other calls (not typically for internal proxy)
# RUNPOD_AI_API_KEY = os.environ.get('RUNPOD_AI_API_KEY')


def handler(job):
    print("---- Runpod Handler: Job Received ----")
    print(f"Runpod Handler: Full job object: {job}")

    job_input = job.get("input")

    if job_input is None:
        print("---- Runpod Handler: Error - 'input' field missing from job ----")
        return {
            "error": "The 'input' field is missing from the job payload."
        }

    print(f"Runpod Handler: Extracted job_input: {job_input}")

    try:
        # Forward the job_input directly to the FastAPI /tts endpoint
        # The job_input should be the JSON body expected by your /tts endpoint
        print(
            f"Runpod Handler: Forwarding request to FastAPI: {FASTAPI_SERVER_URL}")

        # Using httpx for the request
        # Timeout is important for serverless functions
        # 10s total, 5s connect, 60s read
        timeout_config = httpx.Timeout(10.0, connect=5.0, read=60.0)

        response = httpx.post(
            FASTAPI_SERVER_URL,
            json=job_input,  # Send the extracted input as JSON
            timeout=timeout_config
        )

        print(
            f"Runpod Handler: FastAPI response status code: {response.status_code}")

        # Check if the response was successful
        if response.status_code == 200:
            # Assuming the FastAPI server returns audio data directly
            # For Runpod, you might need to decide how to return this.
            # If it's binary data (like a WAV file), you might need to
            # base64 encode it or save it and return a URL if Runpod supports that.
            # For now, let's try to return what FastAPI gives.
            # If FastAPI returns JSON (e.g. with an error or audio URL), that's fine.

            content_type = response.headers.get("content-type", "").lower()

            if "application/json" in content_type:
                print("Runpod Handler: FastAPI returned JSON.")
                return response.json()
            elif "audio" in content_type:
                print(
                    f"Runpod Handler: FastAPI returned audio ({content_type}). Returning raw bytes might not be directly supported by default for JSON return. This part may need adjustment based on how Runpod expects binary data.")
                # This is tricky. Runpod expects a JSON serializable output.
                # Returning raw audio bytes directly might not work as the top-level return.
                # You typically return a JSON with a link to the audio or base64 encoded audio.
                # For initial testing, let's see if Runpod handles it or what error it gives.
                # A more robust way is to save the audio and return a (temp) URL or base64.
                # However, since our FastAPI already returns StreamingResponse for audio,
                # the raw bytes might be what httpx gets.
                # For now, let's just return a success message if it's audio, as the actual audio
                # is streamed by FastAPI. The handler's role is to trigger it.
                # The Runpod platform itself might handle the stream proxying if the Content-Type is correct.

                # If FastAPI streams the audio, the handler's job is done once the request is made.
                # What we return here is more for the Runpod job status.
                print(
                    f"Runpod Handler: Audio response received from FastAPI. Content-Type: {content_type}")
                # Let's assume for now that if we get a 200 with audio, the stream is being handled,
                # and we just return a success indicator for the job.
                # The actual audio would be streamed back to the original Runpod /runsync or /stream caller.
                return {"message": "Audio generation request successful. Audio is being streamed.", "content_type": content_type}

            else:
                print(
                    f"Runpod Handler: FastAPI returned non-JSON, non-audio content-type: {content_type}. Returning as text.")
                return {"response_text": response.text, "status_code": response.status_code}

        else:
            # If FastAPI returned an error
            print(
                f"Runpod Handler: Error from FastAPI - Status: {response.status_code}, Body: {response.text}")
            error_detail = "Error from TTS service."
            try:
                error_detail = response.json()  # Try to parse FastAPI's error JSON
            except ValueError:
                error_detail = response.text  # Fallback to raw text
            return {
                "error": error_detail,
                "status_code_from_fastapi": response.status_code
            }

    except httpx.RequestError as exc:
        print(
            f"---- Runpod Handler: HTTPX RequestError connecting to FastAPI: {exc} ----")
        return {
            "error": f"Failed to connect to the TTS service: {str(exc)}"
        }
    except Exception as e:
        print(f"---- Runpod Handler: Unexpected Error in handler: {e} ----")
        import traceback
        traceback.print_exc()
        return {
            "error": f"An unexpected error occurred in the handler: {str(e)}"
        }


if __name__ == '__main__':
    print("---- Runpod Handler: Starting FastAPI server in background... ----")
    try:
        fastapi_process = subprocess.Popen(
            [sys.executable, "server.py"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        print(
            f"---- Runpod Handler: FastAPI server process started with PID: {fastapi_process.pid} ----")
        print(
            "---- Runpod Handler: Giving FastAPI server a few seconds to initialize... ----")
        # Increased sleep to give ample time for model loading in FastAPI
        time.sleep(10)
        print("---- Runpod Handler: FastAPI server should be initializing. ----")

    except Exception as e_proc:
        print(
            f"---- Runpod Handler: FAILED to start FastAPI server: {e_proc} ----")

    print("---- Runpod Handler: Starting serverless worker (runpod.serverless.start) ----")
    runpod.serverless.start({"handler": handler})