from flask import Flask, jsonify, request, send_file, send_from_directory
from PIL import Image, ImageDraw, ImageFont
import os
import io
import hashlib
import time
from functools import wraps

app = Flask(__name__, static_folder='static')

# API Key authentication
API_KEY = os.environ.get("HTMAA_API_KEY")
if not API_KEY:
    raise ValueError("HTMAA_API_KEY environment variable must be set")

# Storage for images and current selection
IMAGES_DIR = "images"
CURRENT_IMAGE_FILE = os.path.join(IMAGES_DIR, ".current")
CURRENT_HASH_FILE = os.path.join(IMAGES_DIR, ".current_hash")
os.makedirs(IMAGES_DIR, exist_ok=True)


def get_current_image_name():
    """Read the current image name from file."""
    if os.path.exists(CURRENT_IMAGE_FILE):
        with open(CURRENT_IMAGE_FILE, "r") as f:
            return f.read().strip() or None
    return None


def get_current_image_hash():
    """Read the hash of the current image."""
    if os.path.exists(CURRENT_HASH_FILE):
        with open(CURRENT_HASH_FILE, "r") as f:
            return f.read().strip() or None
    return None


def compute_image_hash(image_path):
    """Compute SHA256 hash of an image file."""
    sha256_hash = hashlib.sha256()
    with open(image_path, "rb") as f:
        sha256_hash.update(f.read())
    return sha256_hash.hexdigest()


def set_current_image_name(name):
    """Write the current image name and hash to file."""
    with open(CURRENT_IMAGE_FILE, "w") as f:
        f.write(name if name else "")
    
    # Compute and record hash
    if name:
        image_path = os.path.join(IMAGES_DIR, f"{name}.bin")
        if os.path.exists(image_path):
            image_hash = compute_image_hash(image_path)
            with open(CURRENT_HASH_FILE, "w") as f:
                f.write(image_hash)
        else:
            with open(CURRENT_HASH_FILE, "w") as f:
                f.write("")
    else:
        with open(CURRENT_HASH_FILE, "w") as f:
            f.write("")


def require_api_key(f):
    """Decorator to require API key authentication."""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        # Let CORS preflight through without auth
        if request.method == 'OPTIONS':
            return '', 200

        auth_header = request.headers.get('X-API-Key')
        if not auth_header or auth_header != API_KEY:
            return jsonify(ok=False, error="Invalid or missing API key"), 401
        return f(*args, **kwargs)
    return decorated_function


def pack_1bpp(im, invert=False, msb_first=True):
    """
    Convert PIL image to 1-bit packed binary format.
    Returns (bytes, width, height).
    """
    # Convert to 1-bit (black and white)
    im = im.convert("1")
    w, h = im.size
    pixels = list(im.getdata())
    
    # Pack pixels into bytes
    buf = bytearray()
    for y in range(h):
        for x in range(0, w, 8):
            byte = 0
            for bit in range(8):
                if x + bit < w:
                    pixel_idx = y * w + x + bit
                    pixel = pixels[pixel_idx]
                    # In mode '1', 0 is black, 255 is white
                    bit_value = 0 if pixel == 0 else 1
                    if invert:
                        bit_value = 1 - bit_value
                    
                    if msb_first:
                        byte |= (bit_value << (7 - bit))
                    else:
                        byte |= (bit_value << bit)
            buf.append(byte)
    
    return bytes(buf), w, h


@app.get("/healthz")
def health():
    return jsonify(ok=True)


@app.get("/")
def index():
    """Serve the main HTML viewer page."""
    return send_from_directory(app.static_folder, 'index.html')


@app.get("/current")
@require_api_key
def get_current():
    """
    Get the current image name and hash.
    Returns: {'ok': true, 'name': 'image_name', 'hash': 'abc123...'}
    """
    current_name = get_current_image_name()
    current_hash = get_current_image_hash()
    
    if current_name is None:
        return jsonify(ok=True, name=None, hash=None)
    
    return jsonify(ok=True, name=current_name, hash=current_hash)


@app.get("/images")
@require_api_key
def images_list():
    """
    List all available image names/keys.
    """
    try:
        files = os.listdir(IMAGES_DIR)
        # Strip .bin extension to get image names
        image_names = [f[:-4] for f in files if f.endswith('.bin')]
        return jsonify(ok=True, images=image_names)
    except Exception as e:
        return jsonify(ok=False, error=str(e)), 500


@app.post("/current")
@require_api_key
def set_current():
    """
    Set the current image by name.
    Expected JSON: {'name': 'image_name'}
    """
    data = request.get_json()
    if not data or 'name' not in data:
        return jsonify(ok=False, error="No name provided"), 400
    
    image_name = data['name']
    image_path = os.path.join(IMAGES_DIR, f"{image_name}.bin")
    
    if not os.path.exists(image_path):
        return jsonify(ok=False, error=f"Image '{image_name}' not found"), 404
    
    set_current_image_name(image_name)
    return jsonify(ok=True, current=image_name)


@app.get("/image")
@require_api_key
def image_get():
    """
    Get the current image binary file.
    Query param 'hash' (optional): Client's last known hash.
    If provided and matches current hash, returns JSON indicating no update.
    Otherwise returns the image binary.
    """
    client_hash = request.args.get('hash')
    current_image_name = get_current_image_name()
    
    if current_image_name is None:
        return jsonify(ok=False, error="No current image set"), 404
    
    # If client provided hash, check if update is needed
    current_hash = get_current_image_hash()
    
    if client_hash is not None and client_hash == current_hash:
        return jsonify(ok=True, updated=False, message="No update needed")
    
    image_path = os.path.join(IMAGES_DIR, f"{current_image_name}.bin")
    
    if not os.path.exists(image_path):
        return jsonify(ok=False, error=f"Current image '{current_image_name}' not found"), 404
    
    # Send the image with hash in response header
    response = send_file(image_path, mimetype='application/octet-stream', 
                        as_attachment=True, download_name=f"{current_image_name}.bin")
    if current_hash:
        response.headers['X-Image-Hash'] = current_hash
    return response


@app.get("/image/<name>")
@require_api_key
def image_get_by_name(name):
    """
    Get a specific image binary file by name.
    Returns the raw .bin file.
    """
    image_path = os.path.join(IMAGES_DIR, f"{name}.bin")
    
    if not os.path.exists(image_path):
        return jsonify(ok=False, error=f"Image '{name}' not found"), 404
    
    return send_file(image_path, mimetype='application/octet-stream')


@app.delete("/image/<name>")
@require_api_key
def image_delete(name):
    """
    Delete an image by name.
    """
    image_path = os.path.join(IMAGES_DIR, f"{name}.bin")
    
    if not os.path.exists(image_path):
        return jsonify(ok=False, error=f"Image '{name}' not found"), 404
    
    try:
        # If this is the current image, clear the current selection
        if get_current_image_name() == name:
            set_current_image_name(None)
        
        os.remove(image_path)
        return jsonify(ok=True, deleted=name)
    except Exception as e:
        return jsonify(ok=False, error=str(e)), 500


@app.post("/image")
@require_api_key
def image_post():
    """
    Upload an image with a name. Processes it to 1bpp format and saves to disk.
    Expected form data: 'image' (file), 'name' (string)
    """
    if 'image' not in request.files:
        return jsonify(ok=False, error="No image file provided"), 400
    
    if 'name' not in request.form:
        return jsonify(ok=False, error="No image name provided"), 400
    
    image_file = request.files['image']
    image_name = request.form['name']
    
    if not image_file:
        return jsonify(ok=False, error="Empty image file"), 400
    
    try:
        # Open and process the image
        im = Image.open(image_file)
        
        # Check size and resize if needed
        if im.size != (200, 200):
            print(f"Warning: Image '{image_name}' is {im.size[0]}x{im.size[1]}, resizing to 200x200")
            im = im.resize((200, 200))
        
        # Pack to 1bpp format
        buf, w, h = pack_1bpp(im, invert=False, msb_first=True)
        
        # Sanity check
        expected = (w * h + 7) // 8
        if len(buf) != expected:
            print(f"Warning: got {len(buf)} bytes, expected {expected}")
        
        # Save to disk
        output_path = os.path.join(IMAGES_DIR, f"{image_name}.bin")
        with open(output_path, "wb") as f:
            f.write(buf)
        
        return jsonify(ok=True, name=image_name, width=w, height=h, size=len(buf))
    
    except Exception as e:
        return jsonify(ok=False, error=str(e)), 500


@app.post("/text")
@require_api_key
def text_to_image():
    """
    Convert text to an image and set it as current.
    Expected JSON: {'text': 'your text here', 'x': 10, 'y': 10} (x, y are optional)
    Automatically generates a timestamp-based name and sets as current.
    Text will be wrapped to fit within the 200px width.
    """
    data = request.get_json()
    if not data or 'text' not in data:
        return jsonify(ok=False, error="No text provided"), 400
    
    text = data['text']
    x = data.get('x', 10)
    y = data.get('y', 10)
    font_size = data.get('font_size', 10)
    
    try:
        # Create a white 200x200 image
        im = Image.new('1', (200, 200), 1)  # 1 = white in mode '1'
        draw = ImageDraw.Draw(im)
        
        # Try to load a font with the specified size
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", font_size)
        except:
            print("Warning: Could not load TTF font, using default font")
            font = ImageFont.load_default()
        
        # Word wrap text to fit within width
        max_width = 200 - x - 10  # Leave margin on right
        lines = []
        words = text.split()
        current_line = []
        
        for word in words:
            test_line = ' '.join(current_line + [word])
            bbox = draw.textbbox((0, 0), test_line, font=font)
            line_width = bbox[2] - bbox[0]
            
            if line_width <= max_width:
                current_line.append(word)
            else:
                if current_line:
                    lines.append(' '.join(current_line))
                current_line = [word]
        
        if current_line:
            lines.append(' '.join(current_line))
        
        # Draw wrapped text (0 = black in mode '1')
        wrapped_text = '\n'.join(lines)
        draw.multiline_text((x, y), wrapped_text, fill=0, font=font)
        
        # Pack to 1bpp format
        buf, w, h = pack_1bpp(im, invert=False, msb_first=True)
        
        # Generate timestamp-based name
        image_name = f"text_{int(time.time())}"
        
        # Save to disk
        output_path = os.path.join(IMAGES_DIR, f"{image_name}.bin")
        with open(output_path, "wb") as f:
            f.write(buf)
        
        # Set as current image
        set_current_image_name(image_name)
        
        return jsonify(ok=True, name=image_name, width=w, height=h, size=len(buf), current=True)
    
    except Exception as e:
        return jsonify(ok=False, error=str(e)), 500