import os
import uuid
import boto3
from botocore.exceptions import ClientError
import requests
from IPython.display import display, HTML
import shutil
import pandas as pd
from twelvelabs import TwelveLabs
from qdrant_client import QdrantClient, models
from qdrant_client.models import PointStruct
import time

# AWS S3 Configuration
AWS_ACCESS_KEY = "YOUR_AWS_ACCESS_KEY"
AWS_SECRET_KEY = "YOUR_AWS_SECRET_KEY"
AWS_BUCKET_NAME = "YOUR_BUCKET_NAME"
AWS_REGION = "us-east-1"  # Change to your region

# Initialize S3 client
s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY,
    aws_secret_access_key=AWS_SECRET_KEY,
    region_name=AWS_REGION
)

# Twelve Labs Configuration
TWELVE_LABS_API_KEY = "YOUR_TWELVE_LABS_API_KEY"
twelvelabs_client = TwelveLabs(api_key=TWELVE_LABS_API_KEY)

# Qdrant Configuration
QDRANT_HOST = "YOUR_QDRANT_HOST"
QDRANT_API_KEY = "YOUR_QDRANT_API_KEY"
COLLECTION_NAME = "content_collection"
VECTOR_SIZE = 1024  # Size of embeddings from Twelve Labs

# Initialize Qdrant client
qdrant_client = QdrantClient(
    url=f"https://{QDRANT_HOST}",
    api_key=QDRANT_API_KEY,
    timeout=20,
    prefer_grpc=False
)

# Get a list of video files
video_dir = "downloads/video_content"
video_files = [f for f in os.listdir(video_dir) if f.endswith('.mp4')]

def upload_to_s3(file_path, filename):
    try:
        # Upload the file
        s3_client.upload_file(
            file_path,
            AWS_BUCKET_NAME,
            f"videos-embed/{filename}",
            ExtraArgs={
                'ACL': 'public-read',
                'ContentType': 'video/mp4'
            }
        )

        # Generate the public URL
        url = f"https://{AWS_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/videos-embed/{filename}"
        print(f"Uploaded to S3: {url}")
        return url

    except ClientError as e:
        print(f"Error uploading to S3: {str(e)}")
        raise

def create_video_embedding(video_path, max_retries=3, retry_delay=5):
    if not twelvelabs_client:
        raise ValueError("Twelve Labs API key not configured")

    retries = 0
    while retries < max_retries:
        try:
            print(f"Creating whole video embedding for {video_path}... (Attempt {retries+1}/{max_retries})")

            # Use video_embedding_scopes parameter set to ["clip", "video"] to get whole video embedding
            task = twelvelabs_client.embed.task.create(
                model_name="Marengo-retrieval-2.7",
                video_file=video_path,
                video_embedding_scopes=["clip", "video"]
            )

            print(f"Created task: id={task.id}, status={task.status}")
            task.wait_for_done(sleep_interval=3)
            task_result = twelvelabs_client.embed.task.retrieve(task.id)

            if task_result.status != 'ready':
                raise ValueError(f"Task failed with status: {task_result.status}")

            return task_result

        except Exception as e:
            print(f"Error creating embedding (attempt {retries+1}): {str(e)}")
            retries += 1
            if retries < max_retries:
                print(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
                retry_delay *= 2
            else:
                print("Max retries reached, giving up.")
                raise

def store_in_qdrant(task_result, video_id, s3_url, original_filename):
    if not qdrant_client:
        raise ValueError("Qdrant client not configured")

    try:
        print(f"Processing video embedding for {video_id}...")

        # The embedding will be in the segments with embedding_scope="video"
        if task_result.video_embedding and task_result.video_embedding.segments:
            video_segments = [s for s in task_result.video_embedding.segments
                             if hasattr(s, 'embedding_scope') and s.embedding_scope == 'video']

            if video_segments:
                print(f"Found video-scope embedding")
                embedding_vector = video_segments[0].embeddings_float
            else:
                # If no video scope segment is found, use the first segment as fallback
                print(f"No video-scope embedding found, using first available segment")
                embedding_vector = task_result.video_embedding.segments[0].embeddings_float
        else:
            raise ValueError("No embeddings found in the response")

        # Create a unique point structure for Qdrant storage
        point = PointStruct(
            id=uuid.uuid4().int & ((1<<64)-1), # Generate a unique 64-bit integer ID
            vector=embedding_vector, # Store the extracted embedding vector
            payload={
                'video_id': video_id,
                'video_url': s3_url,  # Store the public S3 URL of the video
                'is_url': True,
                'original_filename': original_filename # Save the original filename
            }
        )

	 # Insert the generated embedding point into the Qdrant collection
        qdrant_client.upsert(collection_name=COLLECTION_NAME, points=[point])
        print(f"Stored whole video embedding in Qdrant")
        return 1
    except Exception as e:
        print(f"Error storing in Qdrant: {str(e)}")
        raise

# Process each video
for filename in video_files[:5]:  # Process first 5 videos or you can setup as per convenience
    try:
        print(f"\nProcessing {filename}...")
        video_path = os.path.join(video_dir, filename)
        video_id = f"{str(uuid.uuid4())[:8]}_{filename}"
        
        # Upload to S3
        s3_url = upload_to_s3(video_path, video_id)
        
        # Generate embeddings
        task_result = create_video_embedding(video_path)
        
        # Store in Qdrant
        store_in_qdrant(task_result, video_id, s3_url, filename)
        
        print(f"Successfully processed {filename}")
    except Exception as e:
        print(f"Error processing {filename}: {str(e)}")

app = Flask(__name__) 
CORS(app, resources={r"/*": {"origins": "*"}})

# Get credentials from environment variables
API_KEY = os.getenv('API_KEY')
QDRANT_HOST = os.getenv('QDRANT_HOST')
QDRANT_API_KEY = os.getenv('QDRANT_API_KEY')

# Qdrant Configuration
COLLECTION_NAME = "content_collection"
VECTOR_SIZE = 1024 # Dimension of vector embeddings

# Initialize clients
try:
    client = TwelveLabs(api_key=API_KEY)
    qdrant_client = QdrantClient(
        url=f"https://{QDRANT_HOST}",
        api_key=QDRANT_API_KEY,
        timeout=20
    )
    logger.info("Successfully initialized API clients")
except Exception as e:
    logger.error(f"Failed to initialize clients: {str(e)}")
    raise

def init_qdrant():
    try:
	 # Fetch all existing collections
        collections = qdrant_client.get_collections().collections
        collection_exists = any(col.name == COLLECTION_NAME for col in collections)
        if not collection_exists:
            # Create the collection with specified vector configuration if it doesn't exist
            qdrant_client.recreate_collection(
                collection_name=COLLECTION_NAME,
                vectors_config=VectorParams(
                    size=VECTOR_SIZE,
                    distance=Distance.COSINE # Use cosine similarity for retrieval
                )
            )
            logger.info(f"Created collection: {COLLECTION_NAME}")
    except Exception as e:
        logger.error(f"Qdrant initialization error: {str(e)}")
        raise

@app.route('/search', methods=['POST'])
def search():
    # Ensure the request contains JSON data
    if not request.is_json:
        logger.warning("Missing JSON data")
        return jsonify({'error': 'Request must be JSON format'}), 400
        
    # Get and validate query
    data = request.get_json()
    query = data.get('query')
    if not query:
        logger.warning("Empty query parameter")
        return jsonify({'error': 'Missing query parameter'}), 400
    
    logger.info(f"Processing search: '{query}'")
    
    try:
        # Generate embedding for the search query
        formatted_query = f"Recommend - {query}"
        embedding_response = client.embed.create(
            model_name="Marengo-retrieval-2.7",
            text=formatted_query
        )

	 # Get the embedding vector
        vector = embedding_response.text_embedding.segments[0].embeddings_float
        
        # Similarity search from the Qdrant collection
        query_response = qdrant_client.query_points(
            collection_name=COLLECTION_NAME,
            query=vector,
            limit=10,
            with_payload=True
        )
        
        # Extract and format results
        search_results = query_response.points
        logger.info(f"Found {len(search_results)} matching results")
        
        # If no results, return empty list
        if not search_results:
            return jsonify([])
        
        # Build formatted response
        formatted_results = []
        for result in search_results:
            point_id = result.id
            score = float(result.score)
            payload = result.payload
            
            formatted_results.append({
                'video_id': payload.get('video_id', f"video_{point_id}"),
                'filename': payload.get('original_filename', payload.get('filename', 'video.mp4')),
                'start_time': float(payload.get('start_time', 0)),
                'end_time': float(payload.get('end_time', 30)),
                'score': score,
                'confidence': 'high' if score > 0.7 else 'medium',
                'url': payload.get('video_url')
            })
        
        logger.info(f"Returning {len(formatted_results)} results")
        return jsonify(formatted_results)
        
    except Exception as e:
        logger.exception(f"Search error: {str(e)}")
        return jsonify({'error': 'Search failed', 'details': str(e)}), 500