import os
import io
import zipfile
from flask import Flask, render_template, request, send_file, jsonify
from werkzeug.exceptions import RequestEntityTooLarge
from PIL import Image, ImageOps

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max upload

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

@app.route('/api/health')
def health():
    return jsonify({'status': 'ok'}), 200

@app.route('/api/split', methods=['POST'])
def split_image():
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    try:
        # Get parameters
        rows = int(request.form.get('rows', 3))
        cols = int(request.form.get('cols', 3))
        out_format = (request.form.get('format') or '').upper()
        quality = int(request.form.get('quality', 90))
        square = request.form.get('square', 'false') == 'true'
        strict_equal = request.form.get('strict_equal', 'false') == 'true'
        
        # Load image
        img = Image.open(file.stream)
        # Auto-orient by EXIF if present
        try:
            img = ImageOps.exif_transpose(img)
        except Exception:
            pass
        
        # Convert to RGB if necessary (e.g. for RGBA/P images saving as JPEG)
        # But we'll try to keep original format if possible, or default to PNG/JPG
        format = img.format or 'PNG'
        if out_format in ['JPEG', 'PNG', 'WEBP']:
            format = out_format
        if format not in ['JPEG', 'PNG', 'WEBP']:
            format = 'PNG'
            
        # Calculate dimensions
        width, height = img.size
        # Optional square crop (center)
        if square:
            side = min(width, height)
            left = (width - side) // 2
            top = (height - side) // 2
            img = img.crop((left, top, left + side, top + side))
            width, height = img.size

        # Compute base tile size
        tile_width = width // cols
        tile_height = height // rows
        if strict_equal:
            # Crop to exact multiple to make all tiles equal
            exact_w = tile_width * cols
            exact_h = tile_height * rows
            left = (width - exact_w) // 2
            top = (height - exact_h) // 2
            img = img.crop((left, top, left + exact_w, top + exact_h))
            width, height = img.size
            tile_width = width // cols
            tile_height = height // rows
        
        # Create ZIP in memory
        memory_file = io.BytesIO()
        with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
            for r in range(rows):
                for c in range(cols):
                    # Crop
                    left = c * tile_width
                    upper = r * tile_height
                    right = left + tile_width
                    lower = upper + tile_height
                    
                    # Handle last row/col to include remainder pixels when not strict_equal
                    if not strict_equal:
                        if c == cols - 1:
                            right = width
                        if r == rows - 1:
                            lower = height
                        
                    tile = img.crop((left, upper, right, lower))
                    
                    # Save tile to bytes
                    tile_bytes = io.BytesIO()
                    save_kwargs = {}
                    if format == 'JPEG':
                        if tile.mode in ('RGBA', 'P'):
                            tile = tile.convert('RGB')
                        save_kwargs['quality'] = max(1, min(quality, 95))
                        save_kwargs['optimize'] = True
                        save_kwargs['progressive'] = True
                    elif format == 'WEBP':
                        save_kwargs['quality'] = max(1, min(quality, 95))
                    tile.save(tile_bytes, format=format, **save_kwargs)
                    tile_bytes.seek(0)
                    
                    # Add to zip (1-based index)
                    # Naming convention: {original_name}_{index}.{ext}
                    # Grid order: 1, 2, 3...
                    index = r * cols + c + 1
                    filename = f"tile_{index:02d}.{format.lower()}"
                    zf.writestr(filename, tile_bytes.read())
            
            # Add posting order guide
            order_lines = []
            order_lines.append(f"发布顺序指南（{rows}x{cols}）")
            order_lines.append("从左到右，从上到下：")
            nums = [f"{i:02d}" for i in range(1, rows * cols + 1)]
            for r in range(rows):
                row_slice = nums[r * cols:(r + 1) * cols]
                order_lines.append("  " + " ".join(row_slice))
            order_lines.append("")
            order_lines.append("建议：")
            order_lines.append("1. 朋友圈/小红书九宫格：按上述顺序依次选择图片。")
            order_lines.append("2. 若平台有排序规则，请关闭自动排序或手动调整。")
            zf.writestr("posting_order.txt", "\n".join(order_lines))
        
        memory_file.seek(0)
        return send_file(
            memory_file,
            mimetype='application/zip',
            as_attachment=True,
            download_name=f'grid_split_{rows}x{cols}.zip'
        )

    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.errorhandler(RequestEntityTooLarge)
def handle_file_too_large(e):
    return jsonify({'error': '文件过大，超过限制（最大16MB）'}), 413

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)