from babyagi.functionz.core.framework import func

@func.register_function(
  metadata={"description": "Generate parameters for Augie creation using GPT."},
  dependencies=["gpt_call"]
)
def generate_augie_params(user_input, voice_id="29vD33N1CtxCmqQRPOHJ"):
  """
  This function generates JSON parameters for creating an Augie video.
  It uses GPT to structure the user input into the required format, keeping the default voice_id.

  Parameters:
  - user_input: The basic user input text.
  - voice_id: Default voice ID (not generated by GPT).

  Returns: A dictionary with the necessary parameters for creating the Augie video.
  """
  prompt = (
      "You are creating parameters for a video generation API request. "
      "The user has provided input text for the video content. Structure the input into the following JSON format:\n"
      "{\n"
      "  'name': '<brief title>',\n"
      "  'text': '<full video content text>',\n"
      "  'orientation': 'landscape' or 'portrait',\n"
      "  'make_video': true or false\n"
      "}\n"
      "Do not generate a voice ID, use the one provided by the API system."
  )

  gpt_output = gpt_call({"prompt": prompt, "user_input": user_input})

  # Parse GPT output and construct parameters
  params = gpt_output['text']  # Assuming gpt_call returns a structured response.
  params['voice_id'] = voice_id  # Set the default voice ID.

  return params


@func.register_function(
  metadata={"description": "Creates a video on Augie platform."},
  key_dependencies=["augie_api_key"],
  imports={"name": "requests", "lib": "requests"}
)
def create_augie(params):
  """Function to create a video on Augie platform with parameters."""
  API_KEY = globals()['augie_api_key']
  BASE_URL = 'https://beta.api.augie.studio/v1'

  headers = {
      'x-api-key': API_KEY,
      'Content-Type': 'application/json'
  }

  import requests
  response = requests.post(f'{BASE_URL}/augies', json=params, headers=headers)

  if response.status_code == 201:
      return response.json()  # Returns the creation response
  else:
      raise Exception(f"Failed to create Augie: {response.text}")


@func.register_function(
  metadata={"description": "Checks the status of the created video."},
  key_dependencies=["augie_api_key"],
  imports={"name": "requests", "lib": "requests"}
)
def get_augie_status(augie_id):
  """Function to check the status of an Augie video creation."""
  API_KEY = globals()['augie_api_key']
  BASE_URL = 'https://beta.api.augie.studio/v1'

  headers = {
      'x-api-key': API_KEY
  }

  import requests
  response = requests.get(f'{BASE_URL}/augies/{augie_id}/status', headers=headers)

  if response.status_code == 200:
      status_data = response.json()
      if status_data.get('status') == 'succeeded' and 'output' in status_data and 'video' in status_data['output']:
          return {"status": "completed", "video_url": status_data['output']['video']}
      else:
          return {"status": "processing"}
  else:
      raise Exception(f"Failed to get Augie status: {response.text}")


@func.register_function(
  metadata={"description": "Wrapper to create a video and keep checking its status until available."},
  dependencies=["generate_augie_params", "create_augie", "get_augie_status"],
  imports={"name": "time", "lib": "time"}
)
def create_and_wait_for_video(user_input, timeout=300, interval=10):
  """
  Wrapper function to create a video from user input and wait for it to be available.
  - user_input: Basic input from the user, processed by GPT.
  - timeout: The max time to wait (in seconds).
  - interval: Time between status checks (in seconds).
  """
  import time

  # Generate parameters using GPT
  params = generate_augie_params(user_input)

  # Create the video
  creation_response = create_augie(params)
  augie_id = creation_response.get('id')

  if not augie_id:
      raise Exception("Failed to retrieve Augie ID after creation.")

  # Wait and check the status periodically
  start_time = time.time()
  while time.time() - start_time < timeout:
      status_response = get_augie_status(augie_id)
      if status_response['status'] == 'completed':
          return status_response  # Return video URL if available
      time.sleep(interval)

  # Timeout reached, return failure
  raise TimeoutError(f"Video creation timed out after {timeout} seconds.")
