frontend/streamlit_app.py

import os
import sys
import time
import json
import tempfile
from pathlib import Path
from datetime import datetime
from typing import Optional

import streamlit as st
import requests
from PIL import Image, ImageDraw, ImageFont
import numpy as np

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Backend API URL
BACKEND_URL = os.environ.get("BACKEND_URL", "http://localhost:8081")

# Page config
st.set_page_config(
    page_title="LEGION Video Generation",
    page_icon="⚔️",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Custom CSS
st.markdown("""
<style>
    /* Main container */
    .main > div {
        padding: 1rem 2rem;
    }
    
    /* Headers */
    h1 {
        font-family: 'Courier New', monospace;
        text-transform: uppercase;
        letter-spacing: 3px;
        color: #ff6b35;
        border-bottom: 2px solid #2a2a3e;
        padding-bottom: 12px;
    }
    
    h2 {
        font-family: 'Courier New', monospace;
        text-transform: uppercase;
        letter-spacing: 2px;
        color: #e0e0e0;
    }
    
    h3 {
        font-family: 'Courier New', monospace;
        color: #ff6b35;
    }
    
    /* Cards */
    .card {
        background: #14141f;
        border: 1px solid #2a2a3e;
        border-radius: 12px;
        padding: 20px;
        margin-bottom: 20px;
    }
    
    /* Buttons */
    .stButton > button {
        background: linear-gradient(135deg, #ff6b35, #e55a2b) !important;
        border: none !important;
        color: white !important;
        font-weight: bold !important;
        letter-spacing: 1px !important;
        width: 100%;
        border-radius: 8px !important;
        padding: 12px 24px !important;
        transition: all 0.3s ease !important;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 5px 20px rgba(255, 107, 53, 0.4);
    }
    
    /* Input fields */
    .stTextInput > div > div > input,
    .stTextArea > div > div > textarea {
        background: #1a1a2e !important;
        border: 1px solid #2a2a3e !important;
        border-radius: 8px !important;
        color: #e0e0e0 !important;
    }
    
    /* Sliders */
    .stSlider > div > div > div > div {
        background: #ff6b35 !important;
    }
    
    /* Status */
    .status-badge {
        background: #00d4aa;
        color: black;
        padding: 4px 12px;
        border-radius: 20px;
        font-size: 0.8em;
        display: inline-block;
    }
    
    /* Footer */
    .footer {
        color: #8888aa;
        border-top: 1px solid #2a2a3e;
        margin-top: 40px;
        padding-top: 15px;
        text-align: center;
    }
    
    /* Video container */
    .stVideo {
        border-radius: 12px;
        overflow: hidden;
    }
    
    /* Tabs */
    .stTabs [data-baseweb="tab-list"] {
        background: #14141f;
        border-radius: 12px 12px 0 0;
        border: 1px solid #2a2a3e;
    }
    
    .stTabs [data-baseweb="tab"] {
        color: #e0e0e0;
        font-family: 'Courier New', monospace;
        letter-spacing: 1px;
    }
    
    .stTabs [aria-selected="true"] {
        background: #ff6b35 !important;
        color: white !important;
    }
    
    /* Sidebar */
    .css-1d391kg {
        background: #0a0a0f;
    }
    
    /* Info boxes */
    .stAlert {
        background: #1a1a2e !important;
        border: 1px solid #2a2a3e !important;
    }
    
    /* Select box */
    .stSelectbox > div > div {
        background: #1a1a2e !important;
        border: 1px solid #2a2a3e !important;
        border-radius: 8px !important;
    }
</style>
""", unsafe_allow_html=True)

# ============================================================
# Sidebar
# ============================================================

with st.sidebar:
    st.markdown("## ⚔️ LEGION")
    st.markdown("### Video Generation Engine")
    st.markdown("---")
    
    # Status check
    try:
        resp = requests.get(f"{BACKEND_URL}/api/status", timeout=3)
        status_data = resp.json()
        status_ok = status_data.get("status") == "ok"
        mock_mode = status_data.get("mock_mode", True)
    except Exception:
        status_ok = False
        mock_mode = True
    
    if status_ok:
        st.success("✅ Backend Connected")
        if mock_mode:
            st.info("🔄 Mock Mode (No GPU)")
        else:
            st.success("🚀 GPU Mode")
    else:
        st.error("❌ Backend Offline")
    
    st.markdown("---")
    st.markdown("### About")
    st.markdown("""
    **LEGION** is a state-of-the-art video generation system with 8.3B parameters.
    
    - Text-to-Video
    - Image-to-Video
    - QWatermark System
    """)
    
    st.markdown("---")
    st.markdown("**v1.0** | Apache 2.0")

# ============================================================
# Main Content
# ============================================================

st.markdown(
    "<h1>⚔️ LEGION VIDEO GENERATION</h1>",
    unsafe_allow_html=True
)

st.markdown(
    '<p style="color: #8888aa; font-size: 1.1em; margin-top: -10px; margin-bottom: 25px;">'
    'The Ultimate AI Video Engine — Text-to-Video & Image-to-Video</p>',
    unsafe_allow_html=True
)

# Tabs
tab1, tab2, tab3 = st.tabs(["🎬 Text-to-Video", "🖼️ Image-to-Video", "💧 QWatermark"])

# ============================================================
# TAB 1: TEXT-TO-VIDEO
# ============================================================

with tab1:
    col1, col2 = st.columns([2, 1])
    
    with col1:
        prompt = st.text_area(
            "📝 Prompt",
            value="A serene mountain lake at sunset with colorful clouds reflecting on the water, gentle ripples, cinematic quality",
            height=120,
            help="Describe the video you want to generate",
        )
        
        negative_prompt = st.text_area(
            "🚫 Negative Prompt",
            value="warped, distorted, flickering, jittery, low quality, blurry, artifacts, ugly, deformed, bad anatomy",
            height=60,
        )
        
        with st.expander("⚙️ Advanced Settings", expanded=False):
            adv_col1, adv_col2 = st.columns(2)
            with adv_col1:
                num_frames = st.slider("Frames", 1, 129, 49, help="Number of frames to generate")
                width = st.selectbox("Width", [256, 384, 480, 720], index=2)
            with adv_col2:
                steps = st.slider("Inference Steps", 10, 100, 50)
                height = st.selectbox("Height", [256, 384, 480, 720], index=2)
            
            col_adv = st.columns(2)
            with col_adv[0]:
                guidance = st.slider("Guidance Scale", 1.0, 20.0, 6.0, 0.5)
            with col_adv[1]:
                wm_strength = st.slider("QWatermark Strength", 0.0, 1.0, 0.3, 0.05)
        
        generate_btn = st.button("🛡️ GENERATE VIDEO", type="primary", use_container_width=True)
    
    with col2:
        video_placeholder = st.empty()
        status_placeholder = st.empty()
    
    if generate_btn:
        if not prompt.strip():
            st.error("Please enter a prompt")
        else:
            with st.spinner("🎬 Generating video... This may take a moment."):
                try:
                    payload = {
                        "prompt": prompt,
                        "negative_prompt": negative_prompt,
                        "num_frames": num_frames,
                        "width": width,
                        "height": height,
                        "num_inference_steps": steps,
                        "guidance_scale": guidance,
                        "watermark_strength": wm_strength,
                    }
                    
                    resp = requests.post(
                        f"{BACKEND_URL}/api/generate/text",
                        json=payload,
                        timeout=600,
                    )
                    
                    if resp.status_code == 200:
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        output_path = f"/tmp/legion_t2v_{timestamp}.mp4"
                        with open(output_path, "wb") as f:
                            f.write(resp.content)
                        
                        video_placeholder.video(output_path)
                        status_placeholder.success("✅ Video generated successfully!")
                    else:
                        status_placeholder.error(f"❌ Error: {resp.text[:200]}")
                except Exception as e:
                    status_placeholder.error(f"❌ Error: {str(e)}")

# ============================================================
# TAB 2: IMAGE-TO-VIDEO
# ============================================================

with tab2:
    col1, col2 = st.columns([2, 1])
    
    with col1:
        uploaded_file = st.file_uploader(
            "📷 Upload Input Image",
            type=["jpg", "jpeg", "png", "webp"],
            help="Upload an image to animate",
        )
        
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            st.image(image, caption="Input Image", use_column_width=True)
            # Save to temp
            temp_dir = tempfile.mkdtemp()
            img_path = os.path.join(temp_dir, uploaded_file.name)
            image.save(img_path)
        
        i2v_prompt = st.text_area(
            "📝 Motion Prompt",
            value="Gentle motion, cinematic camera movement, atmospheric",
            height=80,
            help="Describe the motion or action",
        )
        
        with st.expander("⚙️ Advanced Settings", expanded=False):
            i2v_col1, i2v_col2 = st.columns(2)
            with i2v_col1:
                i2v_frames = st.slider("Frames", 1, 129, 49, key="i2v_frames")
                i2v_width = st.selectbox("Width", [256, 384, 480, 720], index=2, key="i2v_width")
            with i2v_col2:
                i2v_steps = st.slider("Inference Steps", 10, 100, 50, key="i2v_steps")
                i2v_height = st.selectbox("Height", [256, 384, 480, 720], index=2, key="i2v_height")
            
            i2v_guidance = st.slider("Guidance Scale", 1.0, 20.0, 6.0, 0.5, key="i2v_guidance")
            i2v_wm = st.slider("QWatermark Strength", 0.0, 1.0, 0.3, 0.05, key="i2v_wm")
        
        i2v_btn = st.button("🛡️ GENERATE VIDEO", type="primary", use_container_width=True, key="i2v_btn")
    
    with col2:
        i2v_video_placeholder = st.empty()
        i2v_status_placeholder = st.empty()
    
    if i2v_btn:
        if uploaded_file is None:
            st.error("Please upload an image first")
        else:
            with st.spinner("🎬 Generating video from image..."):
                try:
                    with open(img_path, "rb") as f:
                        files = {"file": f}
                        data = {
                            "prompt": i2v_prompt,
                            "negative_prompt": "warped, distorted, flickering, jittery, low quality, blurry",
                            "num_frames": str(i2v_frames),
                            "width": str(i2v_width),
                            "height": str(i2v_height),
                            "num_inference_steps": str(i2v_steps),
                            "guidance_scale": str(i2v_guidance),
                            "watermark_strength": str(i2v_wm),
                        }
                        resp = requests.post(
                            f"{BACKEND_URL}/api/generate/image",
                            files=files,
                            data=data,
                            timeout=600,
                        )
                    
                    if resp.status_code == 200:
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        output_path = f"/tmp/legion_i2v_{timestamp}.mp4"
                        with open(output_path, "wb") as f:
                            f.write(resp.content)
                        
                        i2v_video_placeholder.video(output_path)
                        i2v_status_placeholder.success("✅ Video generated successfully!")
                    else:
                        i2v_status_placeholder.error(f"❌ Error: {resp.text[:200]}")
                except Exception as e:
                    i2v_status_placeholder.error(f"❌ Error: {str(e)}")

# ============================================================
# TAB 3: QWATERMARK SETTINGS
# ============================================================

with tab3:
    st.markdown("## 💧 QWatermark Quality Assurance System")
    st.markdown("""
    The QWatermark system imprints a semi-transparent quality assurance marker on every generated video.
    Configure the watermark appearance below.
    """)
    
    wm_col1, wm_col2 = st.columns(2)
    
    with wm_col1:
        wm_text = st.text_input("Watermark Text", value="LEGION")
        wm_position = st.selectbox(
            "Position",
            ["bottom-right", "bottom-left", "top-right", "top-left", "center"],
            index=0,
        )
    
    with wm_col2:
        wm_font_size = st.slider("Font Size", 16, 72, 36, 2)
        wm_opacity = st.slider("Opacity", 0.0, 1.0, 0.3, 0.05)
    
    if st.button("👁️ PREVIEW WATERMARK", use_container_width=True):
        try:
            frame = Image.new("RGB", (480, 480), (20, 20, 35))
            draw = ImageDraw.Draw(frame)
            
            for i in range(0, 480, 20):
                draw.line([(i, 0), (i, 480)], fill=(30, 30, 50), width=1)
                draw.line([(0, i), (480, i)], fill=(30, 30, 50), width=1)
            
            try:
                font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", wm_font_size)
            except:
                font = ImageFont.load_default()
            
            w, h = frame.size
            bbox = draw.textbbox((0, 0), wm_text, font=font)
            text_w = bbox[2] - bbox[0]
            text_h = bbox[3] - bbox[1]
            
            margin = 15
            pos_map = {
                "top-left": (margin, margin),
                "top-right": (w - text_w - margin, margin),
                "bottom-left": (margin, h - text_h - margin),
                "center": ((w - text_w) // 2, (h - text_h) // 2),
                "bottom-right": (w - text_w - margin, h - text_h - margin),
            }
            x, y = pos_map.get(wm_position, pos_map["bottom-right"])
            
            draw.rectangle(
                [x - 10, y - 10, x + text_w + 10, y + text_h + 10],
                fill=(0, 0, 0, int(40 * wm_opacity)),
            )
            draw.text((x, y), wm_text, font=font, fill=(255, 255, 255, int(255 * wm_opacity)))
            
            st.image(frame, caption="Watermark Preview", use_column_width=True)
        except Exception as e:
            st.error(f"Preview error: {e}")
    
    st.markdown("---")
    st.info("""
    **Current Configuration:**
    - Text: LEGION
    - Position: bottom-right
    - System: Semi-transparent overlay with dark background
    
    The QWatermark is applied to every frame during video export.
    Adjust strength in generation tabs (0.0 to disable).
    """)

# ============================================================
# Footer
# ============================================================

st.markdown("---")
st.markdown(
    '<div class="footer">'
    '<p><strong>⚔️ LEGION VIDEO GENERATION</strong> — v1.0 | Apache 2.0 License</p>'
    '<p style="font-size: 0.8em;">Text-to-Video · Image-to-Video · QWatermark System</p>'
    '</div>',
    unsafe_allow_html=True,
)

# ============================================================
# Run
# ============================================================

if __name__ == "__main__":
    # Streamlit runs via `streamlit run frontend/streamlit_app.py`
    pass