import streamlit as st
from huggingface_hub import HfApi
import os
import json
from datetime import datetime
import cv2
import random
from PIL import Image
import io

# Initialize the Hugging Face API with the token
api = HfApi(token=os.getenv("HF_API_TOKEN"))

def extract_thumbnail(video_path, thumbnail_path):
    video = cv2.VideoCapture(video_path)
    
    # Get total number of frames
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Choose a random frame
    random_frame = random.randint(0, total_frames - 1)
    
    # Set the frame position
    video.set(cv2.CAP_PROP_POS_FRAMES, random_frame)
    
    # Read the frame
    success, frame = video.read()
    if success:
        cv2.imwrite(thumbnail_path, frame)
    
    video.release()
    return success

def save_custom_thumbnail(thumbnail_file, thumbnail_path):
    img = Image.open(thumbnail_file)
    img.save(thumbnail_path)
    return True

def generate_metadata(video_name, title, description, uploader, file_location, thumbnail_location):
    return {
        "fileName": video_name,
        "title": title,
        "description": description,
        "uploader": uploader,
        "uploadTimestamp": datetime.now().isoformat(),
        "fileLocation": file_location,
        "thumbnailLocation": thumbnail_location,
        "views": 0,
        "likes": 0
    }

def upload_video_to_hf(video_file, video_name, title, description, uploader, custom_thumbnail=None):
    # Create temp paths
    temp_dir = "temp"
    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
    
    video_path = os.path.join(temp_dir, video_name)
    thumbnail_name = f"{os.path.splitext(video_name)[0]}_thumb.jpg"
    thumbnail_path = os.path.join(temp_dir, thumbnail_name)
    json_name = f"{os.path.splitext(video_name)[0]}-index.json"
    json_path = os.path.join(temp_dir, json_name)
    
    # Write the video content to a file
    with open(video_path, "wb") as f:
        f.write(video_file.read())
    
    # Handle thumbnail
    if custom_thumbnail:
        thumbnail_extracted = save_custom_thumbnail(custom_thumbnail, thumbnail_path)
    else:
        thumbnail_extracted = extract_thumbnail(video_path, thumbnail_path)
    
    if not thumbnail_extracted:
        st.error("Failed to process thumbnail")
        return None
    
    # Upload the video
    video_location = f"videos/{video_name}"
    api.upload_file(
        path_or_fileobj=video_path,
        path_in_repo=video_location,
        repo_id="vericudebuget/ok4231",
        repo_type="space",
    )
    
    # Upload the thumbnail
    thumbnail_location = f"thumbnails/{thumbnail_name}"
    api.upload_file(
        path_or_fileobj=thumbnail_path,
        path_in_repo=thumbnail_location,
        repo_id="vericudebuget/ok4231",
        repo_type="space",
    )
    
    # Generate and upload metadata JSON
    metadata = generate_metadata(video_name, title, description, uploader, video_location, thumbnail_location)
    with open(json_path, "w") as f:
        json.dump(metadata, f, indent=2)
    
    api.upload_file(
        path_or_fileobj=json_path,
        path_in_repo=f"metadata/{json_name}",
        repo_id="vericudebuget/ok4231",
        repo_type="space",
    )
    
    # Cleanup temp files
    for file_path in [video_path, thumbnail_path, json_path]:
        if os.path.exists(file_path):
            os.remove(file_path)
    
    return metadata

# Streamlit app interface
st.title("Upload your video")
st.markdown("---")

# File uploader for video input
uploaded_video = st.file_uploader("Choose video file", type=["mp4", "avi", "mov"])

if uploaded_video:
    # Show the video details form
    with st.form("video_details"):
        st.write("Video Details")
        title = st.text_input("Title", placeholder="Enter video title")
        description = st.text_area("Description", placeholder="Enter video description")
        uploader = st.text_input("Uploader Name", placeholder="Enter your name")
        
        # Optional custom thumbnail uploader
        custom_thumbnail = st.file_uploader("Upload custom thumbnail (optional)", type=["jpg", "jpeg", "png"])
        
        # Upload button within the form
        submit_button = st.form_submit_button("Upload Video")
        
        if submit_button:
            if not title or not uploader:
                st.error("Please fill in the title and uploader name.")
            else:
                with st.spinner("Uploading video, generating thumbnail and metadata..."):
                    metadata = upload_video_to_hf(
                        uploaded_video, 
                        uploaded_video.name, 
                        title, 
                        description, 
                        uploader,
                        custom_thumbnail
                    )
                    if metadata:
                        st.success("Upload completed successfully!")
                        st.json(metadata)

else:
    st.info("Please upload a video file to begin.")