# app.py (API-Only Version)

import os
import gradio as gr
import requests
import json
import time
import base64
import google.auth
import google.auth.transport.requests
from huggingface_hub import login

# --- 1. Configuration and Authentication (Unchanged) ---
GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID")
GCP_LOCATION = os.environ.get("GCP_LOCATION")

# --- Authentication and Sanity Checks Block (Unchanged) ---
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    print("Hugging Face token found. Logging in.")
    login(token=hf_token)
else:
    print("WARNING: Hugging Face token ('HF_TOKEN') not found.")

creds_json_str = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")

if not all([GCP_PROJECT_ID, GCP_LOCATION, creds_json_str]):
    missing_secrets = [s for s, v in {"GCP_PROJECT_ID": GCP_PROJECT_ID, "GCP_LOCATION": GCP_LOCATION, "GOOGLE_APPLICATION_CREDENTIALS_JSON": creds_json_str}.items() if not v]
    error_message = f"FATAL: Missing required secrets: {', '.join(missing_secrets)}."
    print(error_message)
    # This initial error will still be raised if the app can't start
    raise RuntimeError(error_message)

print("All required secrets are loaded. Initializing API service.")
MODEL_ID = "veo-3.0-generate-preview"
API_ENDPOINT = f"{GCP_LOCATION}-aiplatform.googleapis.com"
PREDICT_URL = f"https://{API_ENDPOINT}/v1/projects/{GCP_PROJECT_ID}/locations/{GCP_LOCATION}/publishers/google/models/{MODEL_ID}:predictLongRunning"
FETCH_URL = f"https://{API_ENDPOINT}/v1/projects/{GCP_PROJECT_ID}/locations/{GCP_LOCATION}/publishers/google/models/{MODEL_ID}:fetchPredictOperation"

with open("gcp_creds.json", "w") as f: f.write(creds_json_str)
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
credentials, _ = google.auth.load_credentials_from_file("gcp_creds.json", scopes=SCOPES)

def get_access_token():
    auth_req = google.auth.transport.requests.Request()
    credentials.refresh(auth_req)
    return credentials.token

# --- 2. Core Video Generation Logic (Refactored for API) ---
# The function now returns a final JSON object instead of yielding updates.
def generate_video_api(prompt: str):
    if not prompt:
        return {"status": "error", "message": "Prompt cannot be empty."}
    
    try:
        headers = {"Authorization": f"Bearer {get_access_token()}", "Content-Type": "application/json"}
        payload = {"instances": [{"prompt": prompt}], "parameters": {"aspectRatio": "16:9", "sampleCount": 1, "durationSeconds": 8, "personGeneration": "allow_all", "addWatermark": True, "includeRaiReason": True, "generateAudio": True}}
        
        # Submit job
        response = requests.post(PREDICT_URL, headers=headers, json=payload)
        response.raise_for_status()
        operation_name = response.json()["name"]
        print(f"Successfully submitted job. Operation Name: {operation_name}")
        
        # Poll for result
        MAX_POLL_ATTEMPTS = 60
        for i in range(MAX_POLL_ATTEMPTS):
            print(f"Polling (Attempt {i+1}/{MAX_POLL_ATTEMPTS})...")
            time.sleep(10) # Wait before polling
            
            headers["Authorization"] = f"Bearer {get_access_token()}"
            fetch_payload = {"operationName": operation_name}
            poll_response = requests.post(FETCH_URL, headers=headers, json=fetch_payload)
            poll_response.raise_for_status()
            poll_result = poll_response.json()

            if poll_result.get("done"):
                print("Job finished.")
                response_data = poll_result.get("response", {})
                
                # Case 1: Success, video is present
                if "videos" in response_data and response_data["videos"]:
                    video_base64 = response_data["videos"][0]["bytesBase64Encoded"]
                    return {"status": "success", "video_base64": video_base64}
                
                # Case 2: Failure, an error message is present
                error_message = "Video generation failed."
                if "error" in poll_result:
                    error_details = poll_result["error"].get("message", "No details provided.")
                    error_message += f" API Error: {error_details}"
                elif "raiResult" in response_data:
                    rai_reason = response_data.get("raiMediaFilteredReason", "Unknown reason.")
                    error_message += f" Content was blocked by safety filters ({rai_reason})."
                else:
                    error_message += " The API did not return a video or a specific error."
                
                return {"status": "error", "message": error_message}
                
        return {"status": "error", "message": "Operation timed out."}
        
    except requests.exceptions.HTTPError as e:
        print(f"HTTP Error: {e.response.text}")
        return {"status": "error", "message": f"API Error: {e.response.status_code}. Details: {e.response.text}"}
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        return {"status": "error", "message": f"An unexpected error occurred: {str(e)}"}

# --- 3. Gradio API Definition (No UI) ---
# We define the components that make up the API contract, but they are not made visible.
with gr.Blocks() as demo:
    # Define the inputs and outputs for the API
    prompt_input = gr.Textbox(label="prompt", visible=False)
    output_json = gr.JSON(label="result", visible=False)

    # Create the API endpoint named "predict"
    # This will be available at /run/predict
    gr.Interface(
        fn=generate_video_api,
        inputs=prompt_input,
        outputs=output_json,
        api_name="predict"
    )

# The launch() call is still needed to start the web server that listens for API calls.
demo.launch()