# Install necessary libraries and dependencies
!pip install twelvelabs
# Import required modules
import twelvelabs

from twelvelabs import TwelveLabs
from twelvelabs.models.embed import EmbeddingsTask

# Initialize the Twelve Labs client
twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)

def generate_embedding(video_url):
    # Create an embedding task
    task = twelvelabs_client.embed.task.create(
        engine_name="Marengo-retrieval-2.6",
        video_url=video_url
    )
    print(f"Created task: id={task.id} engine_name={task.engine_name} status={task.status}")

    # Define a callback function to monitor task progress
    def on_task_update(task: EmbeddingsTask):
        print(f"  Status={task.status}")

    # Wait for the task to complete
    status = task.wait_for_done(
        sleep_interval=2,
        callback=on_task_update
    )
    print(f"Embedding done: {status}")

    # Retrieve the task result
    task_result = twelvelabs_client.embed.task.retrieve(task.id)

    # Extract and return the embeddings
    embeddings = []
    for v in task_result.video_embeddings:
        embeddings.append({
            'embedding': v.values,
            'start_offset_sec': v.start_offset_sec,
            'end_offset_sec': v.end_offset_sec,
            'embedding_scope': v.embedding_scope
        })

    return embeddings, task_result

# Example usage
video_url = "https://storage.googleapis.com/ad-demos-datasets/videos/Ecommerce%20v2.5.mp4"

# Generate embeddings for the video
embeddings, task_result = generate_embedding(video_url)

print(f"Generated {len(embeddings)} embeddings for the video")
for i, emb in enumerate(embeddings):
    print(f"Embedding {i+1}:")
    print(f"  Scope: {emb['embedding_scope']}")
    print(f"  Time range: {emb['start_offset_sec']} - {emb['end_offset_sec']} seconds")
    print(f"  Embedding vector (first 5 values): {emb['embedding'][:5]}")
    print()

def create_text_embedding(
    twelvelabs_client: TwelveLabs,
    text: str,
    engine_name: str = "Marengo-retrieval-2.6",
    verbose: bool = True
) -> dict:
    """
    Create a text embedding using Twelve Labs Embed API.
    
    Example:
        twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)
        text_embedding = create_text_embedding(twelvelabs_client, "Your text here")
    """
    # Create embedding
    text_embedding = twelvelabs_client.embed.create(
        engine_name=engine_name,
        text=text
    )
    
    # Print information if verbose is True
    if verbose:
        print("Created a text embedding")
        print(f" Engine: {text_embedding.engine_name}")
        print(f" Embedding: {text_embedding.text_embedding}")
    
    return text_embedding

def create_audio_embedding(
    twelvelabs_client: TwelveLabs,
    audio_file: Union[str, Path],
    engine_name: str = "Marengo-retrieval-2.6",
    verbose: bool = True
) -> dict:
    """
    Create an audio embedding using Twelve Labs Embed API.
    
    Example:
        twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)
        audio_embedding = create_audio_embedding(twelvelabs_client, "path/to/audio.mp3")
    """
    # Create embedding
    audio_embedding = twelvelabs_client.embed.create(
        engine_name=engine_name,
        audio_file=audio_file,
    )
    
    # Print information if verbose is True
    if verbose:
        print("Created an audio embedding")
        print(f" Engine: {audio_embedding.engine_name}")
        
        if audio_embedding.audio_embedding.segments:
            print("Segments:")
            for i, segment in enumerate(audio_embedding.audio_embedding.segments, 1):
                print(f" Segment {i}:")
                print(f" Start Offset (sec): {segment.start_offset_sec}")
                if segment.values:
                    print(f" Values: {segment.values[:5]}... (truncated)")
    
    return audio_embedding

def create_image_embedding(
    twelvelabs_client: TwelveLabs,
    image_file: Union[str, Path],
    engine_name: str = "Marengo-retrieval-2.6",
    verbose: bool = True
) -> dict:
    """
    Create an image embedding using Twelve Labs Embed API.
    
    Example:
        twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)
        image_embedding = create_image_embedding(twelvelabs_client, "path/to/image.jpg")
    """
    # Create embedding
    image_embedding = twelvelabs_client.embed.create(
        engine_name=engine_name,
        image_file=image_file
    )
    
    # Print information if verbose is True
    if verbose:
        print("Created an image embedding")
        print(f" Engine: {image_embedding.engine_name}")
        print(f" Embedding: {image_embedding.image_embedding.values[:5]}... (truncated)")
    
    return image_embedding