app_enhanced.py

import spaces  # <--- CRITICAL: MUST BE THE FIRST LINE
import os
import threading
import json
import traceback
import shutil
import cv2
import numpy as np
import torch
from flask import Flask, jsonify, request, send_from_directory, send_file

# ======================================================
# 🚀 ZEROGPU WARMUP (FIXES RUNTIME ERROR)
# ======================================================
@spaces.GPU
def gpu_warmup():
    """Dummy function to trigger ZeroGPU detection at startup."""
    print(f"✅ GPU Warmup: CUDA Available = {torch.cuda.is_available()}")
    return True

# ======================================================
# 🔧 CONFIGURATION
# ======================================================
app = Flask(__name__)

# Persistent storage check
if os.path.exists('/data'):
    BASE_STORAGE_PATH = '/data'
else:
    BASE_STORAGE_PATH = '.'

BASE_USER_DIR = os.path.join(BASE_STORAGE_PATH, "userdata")
os.makedirs(BASE_USER_DIR, exist_ok=True)

# ======================================================
# 🧠 BACKEND LOGIC
# ======================================================

def create_placeholder_image(text, filename, output_dir):
    """Creates a backup image if video fails to read."""
    # Create dark grey image
    img = np.zeros((800, 800, 3), dtype=np.uint8)
    img[:] = (40, 40, 40) 
    
    # Add text
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(img, text, (50, 400), font, 1.5, (200, 200, 200), 3, cv2.LINE_AA)
    
    # Add border
    cv2.rectangle(img, (0,0), (800,800), (100,100,100), 20)
    
    path = os.path.join(output_dir, filename)
    cv2.imwrite(path, img)
    return filename

@spaces.GPU(duration=120)
def generate_comic_gpu(video_path, frames_dir, target_pages):
    """
    Extracts 4 frames per page (2x2 Grid).
    """
    print(f"🚀 Starting GPU generation for {video_path}")
    
    # 1. Setup
    if os.path.exists(frames_dir): shutil.rmtree(frames_dir)
    os.makedirs(frames_dir, exist_ok=True)
    
    cap = cv2.VideoCapture(video_path)
    fps = 25
    total_frames = 0
    duration = 0
    
    if cap.isOpened():
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS) or 25
        duration = total_frames / fps
    else:
        print("❌ Video load failed. Using placeholders.")

    # 2. Calculate frames needed (4 per page)
    panels_per_page = 4
    total_panels_needed = int(target_pages) * panels_per_page
    
    frame_files_ordered = []
    
    # 3. Extract Frames
    if duration > 0 and total_frames > 0:
        times = np.linspace(1, max(1, duration - 1), total_panels_needed)
        
        for i, t in enumerate(times):
            cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
            ret, frame = cap.read()
            fname = f"frame_{i:04d}.png"
            
            if ret and frame is not None:
                # Crop to square-ish for 2x2 grid
                h, w = frame.shape[:2]
                min_dim = min(h, w)
                start_x = (w - min_dim) // 2
                start_y = (h - min_dim) // 2
                # Crop center square
                frame = frame[start_y:start_y+min_dim, start_x:start_x+min_dim]
                
                # Resize for consistency
                frame = cv2.resize(frame, (800, 800))
                
                cv2.imwrite(os.path.join(frames_dir, fname), frame)
                frame_files_ordered.append(fname)
            else:
                create_placeholder_image(f"Error {t:.1f}s", fname, frames_dir)
                frame_files_ordered.append(fname)
        cap.release()
    else:
        # Fallback loop
        for i in range(total_panels_needed):
            fname = f"placeholder_{i}.png"
            create_placeholder_image(f"Panel {i+1}", fname, frames_dir)
            frame_files_ordered.append(fname)

    # 4. Build Page Data
    pages_data = []
    for i in range(int(target_pages)):
        start = i * panels_per_page
        end = start + panels_per_page
        p_frames = frame_files_ordered[start:end]
        
        # Fill missing if any
        while len(p_frames) < 4:
            fname = f"extra_{len(p_frames)}.png"
            create_placeholder_image("Empty", fname, frames_dir)
            p_frames.append(fname)
            
        pg_panels = [{'image': f} for f in p_frames]
        
        pg_bubbles = []
        if i == 0:
            pg_bubbles.append({'dialog': "Drag the RED DOT\nto resize panels!", 'x': '50%', 'y': '50%'})

        pages_data.append({
            'panels': pg_panels,
            'bubbles': pg_bubbles,
            'splitX': '50%', # Center X
            'splitY': '50%'  # Center Y
        })

    return pages_data

class ComicGenHost:
    def __init__(self, sid):
        self.sid = sid
        self.user_dir = os.path.join(BASE_USER_DIR, sid)
        self.video_path = os.path.join(self.user_dir, 'video.mp4')
        self.frames_dir = os.path.join(self.user_dir, 'frames')
        self.output_dir = os.path.join(self.user_dir, 'output')
        
        os.makedirs(self.user_dir, exist_ok=True)
        os.makedirs(self.frames_dir, exist_ok=True)
        os.makedirs(self.output_dir, exist_ok=True)

    def run(self, pages):
        try:
            self.write_status("Generating...", 30)
            data = generate_comic_gpu(self.video_path, self.frames_dir, pages)
            
            with open(os.path.join(self.output_dir, 'data.json'), 'w') as f:
                json.dump(data, f)
            
            self.write_status("Ready", 100)
        except Exception as e:
            traceback.print_exc()
            self.write_status(f"Error: {e}", -1)

    def write_status(self, msg, prog):
        with open(os.path.join(self.output_dir, 'status.json'), 'w') as f:
            json.dump({'message': msg, 'progress': prog}, f)

# ======================================================
# 🌐 FRONTEND
# ======================================================
INDEX_HTML = '''
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>4-Panel Adjustable Comic</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/html-to-image/1.11.11/html-to-image.min.js"></script>
    <link href="https://fonts.googleapis.com/css2?family=Bangers&display=swap" rel="stylesheet">
    <style>
        body { background: #121212; color: #eee; font-family: sans-serif; margin: 0; text-align: center; }
        
        /* UPLOAD SCREEN */
        #upload-view { padding: 50px; }
        .box { background: #1e1e1e; display: inline-block; padding: 40px; border-radius: 10px; box-shadow: 0 4px 15px rgba(0,0,0,0.5); }
        button { background: #e74c3c; border: none; padding: 10px 20px; font-weight: bold; cursor: pointer; border-radius: 4px; font-size: 16px; margin-top: 10px; color: white; }
        button:hover { background: #c0392b; }
        input { padding: 10px; margin: 5px; border-radius: 4px; border: 1px solid #555; background: #333; color: white; }

        /* EDITOR SCREEN */
        #editor-view { display: none; padding: 20px; }
        .comic-container { display: flex; flex-wrap: wrap; justify-content: center; gap: 40px; margin-top: 20px; padding-bottom: 80px; }
        
        /* COMIC PAGE STYLE */
        .comic-page {
            width: 600px; height: 800px;
            background: white;
            border: 4px solid #000;
            position: relative;
            box-shadow: 0 0 20px rgba(0,0,0,0.5);
            user-select: none;
            overflow: hidden; 
        }

        /* 2x2 GRID LOGIC */
        .comic-grid {
            width: 100%; height: 100%;
            position: relative;
            background: #000; /* The "Gap" color */
            --x: 50%;
            --y: 50%;
            --gap: 4px; /* Thickness of divider */
        }

        .panel {
            position: absolute;
            overflow: hidden;
            background: #333;
            box-sizing: border-box;
            border: 2px solid #000;
        }

        .panel img {
            width: 100%; height: 100%;
            object-fit: cover;
            pointer-events: auto; 
        }

        /* DYNAMIC PANEL SIZING */
        /* Top Left */
        .panel:nth-child(1) { left: 0; top: 0; width: calc(var(--x) - var(--gap)/2); height: calc(var(--y) - var(--gap)/2); }
        /* Top Right */
        .panel:nth-child(2) { left: calc(var(--x) + var(--gap)/2); top: 0; width: calc(100% - var(--x) - var(--gap)/2); height: calc(var(--y) - var(--gap)/2); }
        /* Bottom Left */
        .panel:nth-child(3) { left: 0; top: calc(var(--y) + var(--gap)/2); width: calc(var(--x) - var(--gap)/2); height: calc(100% - var(--y) - var(--gap)/2); }
        /* Bottom Right */
        .panel:nth-child(4) { left: calc(var(--x) + var(--gap)/2); top: calc(var(--y) + var(--gap)/2); width: calc(100% - var(--x) - var(--gap)/2); height: calc(100% - var(--y) - var(--gap)/2); }

        /* CENTRAL HANDLE (RED DOT) */
        .grid-handle {
            position: absolute;
            width: 24px; height: 24px;
            background: #e74c3c;
            border: 3px solid white;
            border-radius: 50%;
            left: var(--x); top: var(--y);
            transform: translate(-50%, -50%);
            cursor: move;
            z-index: 999;
            box-shadow: 0 4px 10px rgba(0,0,0,0.5);
            pointer-events: auto;
        }
        .grid-handle:hover { transform: translate(-50%, -50%) scale(1.2); }

        /* BUBBLES */
        .bubble {
            position: absolute;
            background: white; color: black;
            padding: 10px 15px; border-radius: 20px;
            font-family: 'Bangers', cursive;
            letter-spacing: 1px;
            border: 2px solid black;
            z-index: 100; cursor: move;
            transform: translate(-50%, -50%);
            min-width: 50px; text-align: center;
        }
        .bubble:after {
            content: ''; position: absolute;
            bottom: -10px; left: 50%; transform: translateX(-50%);
            border-width: 10px 10px 0; border-style: solid;
            border-color: black transparent transparent transparent;
        }

        /* CONTROLS */
        .toolbar {
            position: fixed; bottom: 20px; left: 50%; transform: translateX(-50%);
            background: #333; padding: 10px 20px; border-radius: 50px;
            display: flex; gap: 15px; z-index: 1000; box-shadow: 0 5px 20px rgba(0,0,0,0.6);
        }
        .toolbar button { margin-top: 0; font-size: 14px; }
    </style>
</head>
<body>

<div id="upload-view">
    <div class="box">
        <h1>🎞️ 4-Panel Comic Maker</h1>
        <p>Upload a video to create a 2x2 Comic Page.</p>
        <input type="file" id="fileIn" accept="video/*"><br>
        <label>Pages:</label> <input type="number" id="pgCount" value="1" min="1" max="5" style="width:50px;">
        <br>
        <button onclick="startUpload()">🚀 Generate</button>
        <p id="status" style="color: #bbb; margin-top:10px;"></p>
    </div>
</div>

<div id="editor-view">
    <h2>Drag the <span style="color:#e74c3c">Red Dot</span> to resize panels!</h2>
    <div class="comic-container" id="container"></div>
    
    <div class="toolbar">
        <button onclick="addBubble()">💬 Add Text</button>
        <button onclick="downloadAll()">💾 Download</button>
        <button style="background:#555" onclick="location.reload()">↺ New</button>
    </div>
</div>

<script>
    let sid = 'S' + Date.now();
    let dragItem = null; 
    let activeEl = null;

    async function startUpload() {
        let f = document.getElementById('fileIn').files[0];
        if(!f) return alert("Select a video.");
        
        let fd = new FormData();
        fd.append('file', f);
        fd.append('pages', document.getElementById('pgCount').value);
        
        document.getElementById('status').innerText = "Uploading & Processing...";
        
        let r = await fetch(`/upload?sid=${sid}`, {method:'POST', body:fd});
        if(r.ok) monitorStatus();
    }

    function monitorStatus() {
        let t = setInterval(async () => {
            let r = await fetch(`/status?sid=${sid}`);
            let d = await r.json();
            document.getElementById('status').innerText = d.message;
            if(d.progress === 100) {
                clearInterval(t);
                loadEditor();
            }
        }, 1000);
    }

    async function loadEditor() {
        document.getElementById('upload-view').style.display='none';
        document.getElementById('editor-view').style.display='block';
        
        let r = await fetch(`/output/data.json?sid=${sid}`);
        let data = await r.json();
        
        let con = document.getElementById('container');
        con.innerHTML = '';
        
        data.forEach((pg, i) => {
            let page = document.createElement('div');
            page.className = 'comic-page';
            
            let grid = document.createElement('div');
            grid.className = 'comic-grid';
            grid.style.setProperty('--x', pg.splitX || '50%');
            grid.style.setProperty('--y', pg.splitY || '50%');
            
            // 4 Panels
            pg.panels.forEach(pan => {
                let div = document.createElement('div');
                div.className = 'panel';
                // Timestamp avoids caching blank images
                div.innerHTML = `<img src="/frames/${pan.image}?sid=${sid}&t=${Date.now()}">`;
                grid.appendChild(div);
            });
            
            // Handle
            let handle = document.createElement('div');
            handle.className = 'grid-handle';
            handle.onmousedown = (e) => {
                e.stopPropagation();
                dragItem = 'handle';
                activeEl = { handle: handle, grid: grid };
            };
            grid.appendChild(handle);
            
            // Bubbles
            if(pg.bubbles) pg.bubbles.forEach(b => createBubble(b.dialog, grid));
            
            page.appendChild(grid);
            con.appendChild(page);
        });
    }

    function createBubble(txt, parent) {
        let b = document.createElement('div');
        b.className = 'bubble';
        b.contentEditable = true;
        b.innerText = txt || "Text";
        b.style.left = '50%'; b.style.top = '50%';
        
        b.onmousedown = (e) => {
            if(e.target !== b) return;
            e.stopPropagation();
            dragItem = 'bubble';
            activeEl = b;
        };
        
        if(!parent) parent = document.querySelector('.comic-grid');
        if(parent) parent.appendChild(b);
    }
    
    window.addBubble = () => createBubble("New Text");

    // DRAGGING LOGIC
    document.addEventListener('mousemove', (e) => {
        if(!dragItem) return;
        
        if(dragItem === 'handle') {
            let rect = activeEl.grid.getBoundingClientRect();
            let x = e.clientX - rect.left;
            let y = e.clientY - rect.top;
            
            // Constraints (10% to 90%)
            let px = Math.max(10, Math.min(90, (x / rect.width) * 100));
            let py = Math.max(10, Math.min(90, (y / rect.height) * 100));
            
            activeEl.grid.style.setProperty('--x', px + '%');
            activeEl.grid.style.setProperty('--y', py + '%');
        }
        else if(dragItem === 'bubble') {
            let rect = activeEl.parentElement.getBoundingClientRect();
            activeEl.style.left = (e.clientX - rect.left) + 'px';
            activeEl.style.top = (e.clientY - rect.top) + 'px';
        }
    });

    document.addEventListener('mouseup', () => {
        dragItem = null;
        activeEl = null;
    });

    window.downloadAll = async () => {
        let pgs = document.querySelectorAll('.comic-page');
        for(let i=0; i<pgs.length; i++) {
            let handles = pgs[i].querySelectorAll('.grid-handle');
            handles.forEach(h => h.style.display = 'none');
            
            let url = await htmlToImage.toPng(pgs[i]);
            let a = document.createElement('a');
            a.download = `comic_page_${i+1}.png`;
            a.href = url;
            a.click();
            
            handles.forEach(h => h.style.display = 'block');
        }
    };
</script>
</body>
</html>
'''

# ======================================================
# 🔌 FLASK ROUTES
# ======================================================
@app.route('/')
def index():
    return INDEX_HTML

@app.route('/upload', methods=['POST'])
def upload():
    sid = request.args.get('sid')
    f = request.files['file']
    pages = request.form.get('pages', 1)
    
    host = ComicGenHost(sid)
    f.save(host.video_path)
    
    threading.Thread(target=host.run, args=(pages,)).start()
    return jsonify({'ok': True})

@app.route('/status')
def status():
    sid = request.args.get('sid')
    p = os.path.join(BASE_USER_DIR, sid, 'output', 'status.json')
    if os.path.exists(p): return send_file(p)
    return jsonify({'progress': 0, 'message': 'Waiting...'})

@app.route('/frames/<path:filename>')
def frames(filename):
    sid = request.args.get('sid')
    resp = send_from_directory(os.path.join(BASE_USER_DIR, sid, 'frames'), filename)
    resp.headers['Cache-Control'] = 'no-store'
    return resp

@app.route('/output/<path:filename>')
def output(filename):
    sid = request.args.get('sid')
    return send_from_directory(os.path.join(BASE_USER_DIR, sid, 'output'), filename)

if __name__ == '__main__':
    # 🚀 EXPLICIT WARMUP CALL TO REGISTER FUNCTION
    try:
        gpu_warmup()
    except Exception as e:
        print(f"⚠️ Warmup ignored (Normal if not on GPU yet): {e}")

    app.run(host='0.0.0.0', port=7860)