from flask import Flask, session, request, redirect, url_for, render_template_string, send_file
import datetime
import os
import secrets
import torch
from PIL import Image, ImageDraw
from transformers import GroundingDinoProcessor
from modeling_grounding_dino import GroundingDinoForObjectDetection
from itertools import cycle
import tempfile
import io

app = Flask(__name__)
app.secret_key = os.environ.get('SECRET_KEY', secrets.token_hex(16))
SECRET_PASSWORD = "VeronaTrento25!"
app.permanent_session_lifetime = datetime.timedelta(hours=24)

# ===== AUTHENTICATION FUNCTIONS =====
def is_authenticated():
    return session.get('authenticated', False)

def require_auth(f):
    def decorated_function(*args, **kwargs):
        if not is_authenticated():
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    decorated_function.__name__ = f.__name__
    return decorated_function

# ===== ML MODEL SETUP =====
DEVICE = "cpu"
model_id = "fushh7/llmdet_swin_tiny_hf"

print(f"[INFO] Using device: {DEVICE}")
print(f"[INFO] Loading model from {model_id}...")

processor = GroundingDinoProcessor.from_pretrained(model_id)
model = GroundingDinoForObjectDetection.from_pretrained(model_id).to(DEVICE)
model.eval()

print("[INFO] Model loaded successfully.")

# Pre-defined palette
BOX_COLORS = [
    "deepskyblue", "red", "lime", "dodgerblue",
    "cyan", "magenta", "yellow", "orange", "chartreuse"
]

# ===== ML FUNCTIONS =====
def save_cropped_images(original_image, boxes, labels, scores):
    saved_paths = []
    for i, (box, label, score) in enumerate(zip(boxes, labels, scores)):
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
            filepath = tmp_file.name
        cropped_img = original_image.crop(box)
        cropped_img.save(filepath)
        saved_paths.append(filepath)
    return saved_paths

def draw_boxes(image, boxes, labels, scores, colors=BOX_COLORS, font_size=16):
    colour_cycle = cycle(colors)
    draw = ImageDraw.Draw(image)
    
    try:
        font = ImageFont.truetype("arial.ttf", size=font_size)
    except:
        font = ImageFont.load_default()
    
    label_to_colour = {}
    
    for box, label, score in zip(boxes, labels, scores):
        colour = label_to_colour.setdefault(label, next(colour_cycle))
        x_min, y_min, x_max, y_max = map(int, box)
        
        draw.rectangle([x_min, y_min, x_max, y_max], outline=colour, width=2)
        text = f"{label} ({score:.3f})"
        text_bbox = draw.textbbox((0, 0), text, font=font)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        
        bg_coords = [x_min, y_min - text_height - 4, x_min + text_width + 4, y_min]
        draw.rectangle(bg_coords, fill=colour)
        draw.text((x_min + 2, y_min - text_height - 2), text, fill="black", font=font)
    
    return image

def resize_image_max_dimension(image, max_size=1024):
    width, height = image.size
    if max(width, height) <= max_size:
        return image
    ratio = max_size / max(width, height)
    new_width = int(width * ratio)
    new_height = int(height * ratio)
    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)

def detect_and_draw(img, text_query, box_threshold=0.14, text_threshold=0.13):
    text_query = text_query.lower()
    img = resize_image_max_dimension(img, max_size=1024)
    
    inputs = processor(images=img, text=text_query, return_tensors="pt").to(DEVICE)
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    results = processor.post_process_grounded_object_detection(
        outputs,
        inputs.input_ids,
        text_threshold=text_threshold,
        target_sizes=[img.size[::-1]]
    )[0]
    
    img_out = img.copy()
    img_out = draw_boxes(
        img_out,
        boxes=results["boxes"].cpu().numpy(),
        labels=results.get("text_labels", results.get("labels", [])),
        scores=results["scores"]
    )
    
    crop_paths = save_cropped_images(
        img,
        boxes=results["boxes"].cpu().numpy(),
        labels=results.get("text_labels", results.get("labels", [])),
        scores=results["scores"]
    )
    
    return img_out, crop_paths

# ===== FLASK ROUTES =====
@app.route('/')
@require_auth
def index():
    return render_template_string('''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Student Finder - Protetto</title>
        <style>
            body { font-family: Arial, sans-serif; max-width: 1200px; margin: 0 auto; padding: 20px; }
            .header { background: #e8f5e8; padding: 20px; border-radius: 10px; margin-bottom: 20px; }
            .content { background: #f5f5f5; padding: 30px; border-radius: 10px; }
            .form-group { margin-bottom: 15px; }
            label { display: block; margin-bottom: 5px; font-weight: bold; }
            input, textarea, select { width: 100%; padding: 8px; border: 1px solid #ddd; border-radius: 4px; }
            button { background: #007bff; color: white; padding: 10px 20px; border: none; border-radius: 4px; cursor: pointer; }
            button:hover { background: #0056b3; }
            .logout { float: right; }
            .results { margin-top: 20px; }
            .gallery { display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 10px; margin-top: 20px; }
            .gallery img { max-width: 100%; height: auto; border: 1px solid #ddd; border-radius: 4px; }
        </style>
    </head>
    <body>
        <div class="header">
            <h1>🎓 Student Finder</h1>
            <p>Carica una foto di classe e trova gli studenti</p>
            <a href="/logout" class="logout">🔓 Logout</a>
            <div style="clear: both;"></div>
        </div>
        
        <div class="content">
            <form method="post" enctype="multipart/form-data" action="/detect">
                <div class="form-group">
                    <label for="image">Immagine:</label>
                    <input type="file" id="image" name="image" accept="image/*" required>
                </div>
                
                <div class="form-group">
                    <label for="text_query">Text Query:</label>
                    <textarea id="text_query" name="text_query" rows="2" required>heads.</textarea>
                    <small>Testo in lowercase, ogni concetto termina con '.' (es. 'heads. faces.')</small>
                </div>
                
                <div class="form-group">
                    <label for="box_threshold">Box Threshold ({{ box_threshold }}):</label>
                    <input type="range" id="box_threshold" name="box_threshold" min="0" max="1" step="0.05" value="0.14">
                </div>
                
                <div class="form-group">
                    <label for="text_threshold">Text Threshold ({{ text_threshold }}):</label>
                    <input type="range" id="text_threshold" name="text_threshold" min="0" max="1" step="0.05" value="0.13">
                </div>
                
                <button type="submit">🔍 Rileva Studenti</button>
            </form>
            
            {% if result_image %}
            <div class="results">
                <h3>Risultati:</h3>
                <img src="data:image/jpeg;base64,{{ result_image }}" alt="Risultato" style="max-width: 100%;">
                
                {% if crops %}
                <h4>Ritagli individuati ({{ crops|length }}):</h4>
                <div class="gallery">
                    {% for crop in crops %}
                    <img src="data:image/jpeg;base64,{{ crop }}" alt="Ritaglio {{ loop.index }}">
                    {% endfor %}
                </div>
                {% endif %}
            </div>
            {% endif %}
        </div>
    </body>
    </html>
    ''', box_threshold=0.14, text_threshold=0.13)

@app.route('/detect', methods=['POST'])
@require_auth
def detect():
    if 'image' not in request.files:
        return redirect(url_for('index'))
    
    image_file = request.files['image']
    if image_file.filename == '':
        return redirect(url_for('index'))
    
    try:
        # Process image
        image = Image.open(image_file.stream).convert('RGB')
        text_query = request.form.get('text_query', 'heads.')
        box_threshold = float(request.form.get('box_threshold', 0.14))
        text_threshold = float(request.form.get('text_threshold', 0.13))
        
        # Run detection
        result_image, crop_paths = detect_and_draw(image, text_query, box_threshold, text_threshold)
        
        # Convert images to base64 for display
        import base64
        
        # Convert result image to base64
        img_buffer = io.BytesIO()
        result_image.save(img_buffer, format='JPEG')
        result_b64 = base64.b64encode(img_buffer.getvalue()).decode()
        
        # Convert crops to base64
        crops_b64 = []
        for crop_path in crop_paths:
            with open(crop_path, 'rb') as f:
                crop_b64 = base64.b64encode(f.read()).decode()
                crops_b64.append(crop_b64)
            # Cleanup temp file
            os.unlink(crop_path)
        
        return render_template_string('''
        <!DOCTYPE html>
        <html>
        <head>
            <title>Risultati - Student Finder</title>
            <style>
                body { font-family: Arial, sans-serif; max-width: 1200px; margin: 0 auto; padding: 20px; }
                .header { background: #e8f5e8; padding: 20px; border-radius: 10px; margin-bottom: 20px; }
                .content { background: #f5f5f5; padding: 30px; border-radius: 10px; }
                .logout { float: right; }
                .gallery { display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 10px; margin-top: 20px; }
                .gallery img { max-width: 100%; height: auto; border: 1px solid #ddd; border-radius: 4px; }
                .back-btn { background: #6c757d; color: white; padding: 10px 20px; border: none; border-radius: 4px; cursor: pointer; text-decoration: none; display: inline-block; margin-bottom: 20px; }
                .back-btn:hover { background: #545b62; }
            </style>
        </head>
        <body>
            <div class="header">
                <h1>🎓 Risultati Student Finder</h1>
                <a href="/logout" class="logout">🔓 Logout</a>
                <div style="clear: both;"></div>
            </div>
            
            <a href="/" class="back-btn">← Nuova Analisi</a>
            
            <div class="content">
                <h3>Immagine con bounding box:</h3>
                <img src="data:image/jpeg;base64,{{ result_image }}" alt="Risultato" style="max-width: 100%; border: 1px solid #ddd; border-radius: 4px;">
                
                {% if crops %}
                <h3>Ritagli individuati ({{ crops|length }}):</h3>
                <div class="gallery">
                    {% for crop in crops %}
                    <img src="data:image/jpeg;base64,{{ crop }}" alt="Ritaglio {{ loop.index }}">
                    {% endfor %}
                </div>
                {% else %}
                <p>Nessun ritaglio individuato.</p>
                {% endif %}
            </div>
        </body>
        </html>
        ''', result_image=result_b64, crops=crops_b64)
        
    except Exception as e:
        return f"Errore durante l'elaborazione: {str(e)}", 500

@app.route('/login', methods=['GET', 'POST'])
def login():
    if is_authenticated():
        return redirect(url_for('index'))
    
    error = None
    if request.method == 'POST':
        if request.form.get('password') == SECRET_PASSWORD:
            session.permanent = True
            session['authenticated'] = True
            return redirect(url_for('index'))
        else:
            error = "❌ Password errata. Riprova."
    
    return render_template_string('''
    <!DOCTYPE html>
    <html>
    <head>
        <title>Login - Student Finder</title>
        <style>
            body { 
                font-family: Arial, sans-serif; 
                max-width: 400px; 
                margin: 100px auto; 
                padding: 20px;
                background: #f5f5f5;
            }
            .login-form { 
                background: white; 
                padding: 30px; 
                border-radius: 10px;
                box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            }
            h2 {
                color: #333;
                text-align: center;
                margin-bottom: 20px;
            }
            input[type="password"] { 
                width: 100%; 
                padding: 12px; 
                margin: 15px 0; 
                border: 1px solid #ddd; 
                border-radius: 5px;
                box-sizing: border-box;
                font-size: 16px;
            }
            button { 
                background: #007bff; 
                color: white; 
                padding: 12px 20px; 
                border: none; 
                border-radius: 5px; 
                cursor: pointer;
                width: 100%;
                font-size: 16px;
            }
            button:hover {
                background: #0056b3;
            }
            .error { 
                color: red; 
                margin-bottom: 15px;
                text-align: center;
                padding: 10px;
                background: #ffe6e6;
                border-radius: 5px;
            }
        </style>
    </head>
    <body>
        <div class="login-form">
            <h2>🔒 Student Finder - Accesso Protetto</h2>
            <p style="text-align: center; color: #666;">Inserisci la password per accedere</p>
            {% if error %}
                <div class="error">{{ error }}</div>
            {% endif %}
            <form method="POST">
                <input type="password" name="password" placeholder="Password" required>
                <button type="submit">🔑 Accedi</button>
            </form>
        </div>
    </body>
    </html>
    ''', error=error)

@app.route('/logout')
def logout():
    session.clear()
    return redirect(url_for('login'))
    

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False)