# prompt: make a gradio app using ffmpeg.probe to display media info

import gradio as gr
import ffmpeg


def get_media_info(file_path):
  """
  Uses ffmpeg to probe the media file and extract information.
  """
  try:
    probe = ffmpeg.probe(file_path)
    video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
    audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)

    if video_stream:
      width = int(video_stream['width'])
      height = int(video_stream['height'])
      frame_rate = video_stream['avg_frame_rate']
      frame_count = int(video_stream['nb_frames'])
    else:
      width, height, frame_rate, frame_count = None, None, None, None

    if audio_stream:
      sample_rate = int(audio_stream['sample_rate'])
      channels = int(audio_stream['channels'])
    else:
      sample_rate, channels = None, None

    return {
        'width': width,
        'height': height,
        'frame_rate': frame_rate,
        'frame_count': frame_count,
        'sample_rate': sample_rate,
        'channels': channels,
    }

  except Exception as e:
    return {'error': str(e)}


iface = gr.Interface(
    fn=get_media_info,
    inputs=gr.inputs.File(label="Upload Media File"),
    outputs=gr.outputs.JSON(label="Media Information"),
    title="Media Info Extractor",
    description="Upload a media file (video or audio) to extract information using ffmpeg.",
)

iface.launch()