import os
import shutil
import cv2
import tempfile
from flask import Flask, render_template, request, jsonify
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import torch
# We keep gdown import for fallback
import gdown

app = Flask(__name__)

# Google Drive ID for the model (Folder)
MODEL_DRIVE_ID = "1ovFKKSc9xc1-DL9AEAi4ZSGDhiBIFt0j" 
model_path = "./blip-model"

def load_model():
    # Attempt to load model
    try:
        print(f"Loading model from {model_path}...")
        processor = BlipProcessor.from_pretrained(model_path)
        model = BlipForConditionalGeneration.from_pretrained(model_path)
        print("Model loaded successfully.")
        return processor, model
    except Exception as e:
        print(f"Failed to load model from {model_path}: {e}")
        print("Model might not be present. Attempting download (Fallback)...")
        # Fallback download logic here if needed, or we rely on Dockerfile
        try:
             # Check if model exists and is not empty (simple check)
            if not os.path.exists(model_path) or not os.listdir(model_path):
                print(f"Downloading from Google Drive...")
                os.makedirs(model_path, exist_ok=True)
                url = f'https://drive.google.com/drive/folders/{MODEL_DRIVE_ID}'
                gdown.download_folder(url, output=model_path, quiet=False, use_cookies=False)
                
                # Retry load
                processor = BlipProcessor.from_pretrained(model_path)
                model = BlipForConditionalGeneration.from_pretrained(model_path)
                print("Model loaded after fallback download.")
                return processor, model
        except Exception as download_error:
             print(f"Fallback download failed: {download_error}")
             return None, None

processor, model = load_model()

# Check for GPU
if model:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    print(f"Using device: {device}")
else:
    print("CRITICAL WARNING: Model not loaded. Application will start but will fail to generate captions.")
    device = "cpu"

@app.route('/')
def index():
    return render_template('index.html')

def generate_caption(raw_image, style):
    inputs = processor(raw_image, return_tensors="pt").to(device)
    
    gen_kwargs = {
        "max_new_tokens": 100,
        "min_length": 20,
        "num_beams": 5
    }
    
    if style == 'detailed':
        gen_kwargs.update({
            "min_length": 40,
            "max_new_tokens": 150,
            "repetition_penalty": 1.2,
            "num_beams": 5
        })
    elif style == 'concise':
        gen_kwargs.update({
            "max_new_tokens": 20,
            "min_length": 5,
            "num_beams": 3
        })
    elif style == 'creative':
        gen_kwargs.update({
            "do_sample": True,
            "top_k": 50,
            "top_p": 0.92,
            "temperature": 0.8,
            "num_beams": 3,
            "repetition_penalty": 1.3,
            "min_length": 25
        })
    elif style == 'hashtags':
        gen_kwargs.update({
            "max_new_tokens": 30,
            "min_length": 5,
            "num_beams": 3
        })

    out = model.generate(**inputs, **gen_kwargs)
    caption = processor.decode(out[0], skip_special_tokens=True)
    
    if style == 'hashtags':
        words = caption.split()
        hashtags = []
        for w in words:
            clean_w = ''.join(e for e in w if e.isalnum())
            if len(clean_w) > 2:
                hashtags.append(f"#{clean_w.lower()}")
        if not hashtags:
            hashtags = ["#image", "#photo"]
        caption = " ".join(hashtags)

    return caption

def process_video(file, style):
    temp_dir = tempfile.mkdtemp()
    temp_video_path = os.path.join(temp_dir, file.filename)
    file.save(temp_video_path)
    
    captions = []
    try:
        cap = cv2.VideoCapture(temp_video_path)
        fps = round(cap.get(cv2.CAP_PROP_FPS))
        if fps == 0:
            fps = 30 # fallback
            
        frame_count = 0
        success = True
        
        while success:
            success, frame = cap.read()
            if not success:
                break
                
            # Extract 1 frame per second
            if frame_count % fps == 0:
                timestamp_sec = frame_count // fps
                
                # Convert BGR to RGB
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil_img = Image.fromarray(frame_rgb)
                
                caption = generate_caption(pil_img, style)
                captions.append({
                    "timestamp": timestamp_sec,
                    "caption": caption
                })
                
            frame_count += 1
            
        cap.release()
    finally:
        # Cleanup
        if os.path.exists(temp_video_path):
            os.remove(temp_video_path)
        os.rmdir(temp_dir)
        
    return jsonify({'captions': captions, 'type': 'video'})

@app.route('/caption', methods=['POST'])
def caption_image():
    if not model:
        return jsonify({'error': 'Model not loaded'}), 500

    file = request.files.get('file') or request.files.get('image')
    if not file or file.filename == '':
        return jsonify({'error': 'No selected file or image provided'}), 400

    try:
        style = request.form.get('style', 'default')
        mimetype = file.mimetype
        
        if mimetype and mimetype.startswith('video/'):
            return process_video(file, style)
        else:
            raw_image = Image.open(file.stream).convert('RGB')
            caption = generate_caption(raw_image, style)
            return jsonify({'caption': caption, 'type': 'image'})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # Run on 0.0.0.0 and port 7860 for Hugging Face Spaces
    app.run(host='0.0.0.0', port=7860)