import streamlit as st
import os
from pathlib import Path
import time
import base64
from PIL import Image
import io
import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.wsgi import WSGIMiddleware
from starlette.responses import JSONResponse

# Set page config
st.set_page_config(page_title="UltraVideoSpace", layout="wide", initial_sidebar_state="collapsed")

# Custom CSS
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;700&display=swap');
    
    body {
        font-family: 'Roboto', sans-serif;
        background-color: #141414;
        color: #ffffff;
    }
    .stApp {
        background-color: #141414;
    }
    .main .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
    h1, h2, h3 {
        color: #e50914;
        font-weight: 700;
    }
    .stButton>button {
        background-color: #e50914;
        color: white;
        border: none;
        border-radius: 4px;
        padding: 12px 24px;
        font-size: 16px;
        font-weight: 700;
        cursor: pointer;
        transition: all 0.3s ease;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .stButton>button:hover {
        background-color: #f40612;
        transform: translateY(-2px);
        box-shadow: 0 6px 8px rgba(0, 0, 0, 0.15);
    }
    .video-grid {
        display: grid;
        grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
        gap: 24px;
        padding: 24px;
    }
    .video-item {
        position: relative;
        overflow: hidden;
        border-radius: 8px;
        transition: all 0.3s ease;
        box-shadow: 0 4px 10px rgba(0, 0, 0, 0.2);
    }
    .video-item:hover {
        transform: scale(1.05);
        box-shadow: 0 6px 14px rgba(0, 0, 0, 0.3);
    }
    .video-thumbnail {
        width: 100%;
        aspect-ratio: 16 / 9;
        object-fit: cover;
    }
    .video-title {
        position: absolute;
        bottom: 0;
        left: 0;
        right: 0;
        background: linear-gradient(to top, rgba(0,0,0,0.8) 0%, rgba(0,0,0,0) 100%);
        padding: 16px;
        font-size: 14px;
        font-weight: 500;
        text-align: center;
    }
    .upload-zone {
        border: 2px dashed #e50914;
        border-radius: 8px;
        padding: 40px;
        text-align: center;
        cursor: pointer;
        transition: all 0.3s ease;
    }
    .upload-zone:hover {
        background-color: rgba(229, 9, 20, 0.1);
    }
    #file-upload-status {
        margin-top: 20px;
        padding: 10px;
        border-radius: 5px;
        background-color: rgba(0, 0, 0, 0.5);
    }
</style>
""", unsafe_allow_html=True)

# FastAPI app
app = FastAPI()

@app.post("/upload_chunk")
async def upload_chunk(file: UploadFile = File(...), chunk_number: int = Form(...), total_chunks: int = Form(...)):
    chunk_data = await file.read()
    save_chunk(chunk_data, file.filename, chunk_number)
    
    if chunk_number == total_chunks - 1:
        reassemble_file(file.filename, total_chunks)
    
    return JSONResponse(content={"message": "Chunk uploaded successfully"})

# Wrap Streamlit app
streamlit_app = WSGIMiddleware(app)

def save_chunk(chunk, filename, chunk_number):
    save_path = Path("uploaded_chunks")
    save_path.mkdir(exist_ok=True)
    file_path = save_path / f"{filename}.part{chunk_number}"
    with file_path.open("wb") as f:
        f.write(chunk)

def reassemble_file(filename, total_chunks):
    save_path = Path("uploaded_videos")
    save_path.mkdir(exist_ok=True)
    file_path = save_path / filename
    
    with file_path.open("wb") as outfile:
        for i in range(total_chunks):
            chunk_path = Path("uploaded_chunks") / f"{filename}.part{i}"
            with chunk_path.open("rb") as infile:
                outfile.write(infile.read())
            os.remove(chunk_path)

def get_uploaded_videos():
    video_dir = Path("uploaded_videos")
    video_dir.mkdir(exist_ok=True)
    return [f for f in video_dir.glob("*") if f.suffix.lower() in ['.mp4', '.mkv', '.avi', '.mov', '.mpeg4']]

def generate_thumbnail(video_path):
    try:
        video = cv2.VideoCapture(str(video_path))
        success, image = video.read()
        if success:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(image)
            image.thumbnail((300, 168))
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            return base64.b64encode(buffered.getvalue()).decode()
    except Exception as e:
        st.warning(f"Could not generate thumbnail for {video_path.name}: {str(e)}")
    return None

def main():
    st.markdown("<h1 style='text-align: center;'>UltraVideoSpace</h1>", unsafe_allow_html=True)

    st.markdown("""
    <div class="upload-zone" id="drop-zone">
        <h3>Drag and drop your video here</h3>
        <p>or click to select files</p>
        <p>Supported formats: MP4, MKV, AVI, MOV, MPEG4</p>
        <p>No file size limit (files will be split if larger than 200MB)</p>
    </div>
    <div id="file-upload-status"></div>
    """, unsafe_allow_html=True)

    # File uploader and chunked upload handling
    st.markdown("""
    <script>
    const dropZone = document.getElementById('drop-zone');
    const statusDiv = document.getElementById('file-upload-status');

    dropZone.addEventListener('dragover', (e) => {
        e.preventDefault();
        dropZone.style.backgroundColor = 'rgba(229, 9, 20, 0.1)';
    });

    dropZone.addEventListener('dragleave', (e) => {
        e.preventDefault();
        dropZone.style.backgroundColor = '';
    });

    dropZone.addEventListener('drop', (e) => {
        e.preventDefault();
        dropZone.style.backgroundColor = '';
        const file = e.dataTransfer.files[0];
        uploadFile(file);
    });

    dropZone.addEventListener('click', () => {
        const input = document.createElement('input');
        input.type = 'file';
        input.accept = '.mp4,.mkv,.avi,.mov,.mpeg4';
        input.onchange = (e) => {
            const file = e.target.files[0];
            uploadFile(file);
        };
        input.click();
    });

    function uploadFile(file) {
        const chunkSize = 200 * 1024 * 1024; // 200MB
        const totalChunks = Math.ceil(file.size / chunkSize);
        let currentChunk = 0;

        statusDiv.innerHTML = `Uploading ${file.name} (0/${totalChunks} chunks)`;

        function uploadNextChunk() {
            const start = currentChunk * chunkSize;
            const end = Math.min(start + chunkSize, file.size);
            const chunk = file.slice(start, end);

            const formData = new FormData();
            formData.append('file', chunk, file.name);
            formData.append('chunk_number', currentChunk);
            formData.append('total_chunks', totalChunks);

            fetch('/upload_chunk', {
                method: 'POST',
                body: formData
            })
            .then(response => response.json())
            .then(data => {
                currentChunk++;
                statusDiv.innerHTML = `Uploading ${file.name} (${currentChunk}/${totalChunks} chunks)`;
                if (currentChunk < totalChunks) {
                    uploadNextChunk();
                } else {
                    statusDiv.innerHTML = `${file.name} uploaded successfully!`;
                    setTimeout(() => {
                        window.location.reload();
                    }, 2000);
                }
            })
            .catch(error => {
                statusDiv.innerHTML = `Error uploading ${file.name}: ${error}`;
            });
        }

        uploadNextChunk();
    }
    </script>
    """, unsafe_allow_html=True)

    # Display uploaded videos
    st.markdown("<h2 style='text-align: center; margin-top: 40px;'>Your Video Collection</h2>", unsafe_allow_html=True)
    videos = get_uploaded_videos()
    
    if not videos:
        st.info("Your video collection is empty. Upload some videos to get started!")
    else:
        video_grid = "<div class='video-grid'>"
        for video in videos:
            thumbnail = generate_thumbnail(video)
            thumbnail_url = f"data:image/png;base64,{thumbnail}" if thumbnail else "https://via.placeholder.com/300x168.png?text=Video+Thumbnail"
            video_grid += f"""
            <div class='video-item'>
                <img class='video-thumbnail' src='{thumbnail_url}' alt='{video.name}'>
                <div class='video-title'>{video.name}</div>
            </div>
            """
        video_grid += "</div>"
        st.markdown(video_grid, unsafe_allow_html=True)

        selected_video = st.selectbox("Select a video to play", [v.name for v in videos])
        if selected_video:
            video_path = f"uploaded_videos/{selected_video}"
            st.video(video_path)

if __name__ == "__main__":
    main()