matrix_app.py

import logging
import os
import secrets
import subprocess
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from threading import Lock, Semaphore, Thread

from dotenv import load_dotenv
from flask import Flask, Response, abort, jsonify, render_template, request, send_from_directory, session, url_for
from flask_caching import Cache
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from werkzeug.middleware.proxy_fix import ProxyFix

from matrix_pipeline import generate_braille_model_from_text, translate_text_to_braille_text

LOG_DIR = Path(__file__).parent / "logs"
LOG_DIR.mkdir(exist_ok=True)
logging.basicConfig(
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(LOG_DIR / "app.log", encoding="utf-8"),
    ]
)

load_dotenv()


def _get_env_str(name: str) -> str | None:
    value = os.getenv(name)
    if value is None or not value.strip():
        return None
    return value.strip()


def _parse_bool_env(name: str, default: bool) -> bool:
    stripped_value = _get_env_str(name)
    if stripped_value is None:
        return default
    return stripped_value.lower() == "true"


def _parse_int_env(name: str, default: int) -> int:
    stripped_value = _get_env_str(name)
    if stripped_value is None:
        return default
    try:
        return int(stripped_value)
    except ValueError:
        return default


app = Flask(__name__)
app.config["SECRET_KEY"] = _get_env_str("MATRIX_SESSION_SIGNING_KEY") or secrets.token_urlsafe(32)
server_name = _get_env_str("MATRIX_SERVER_NAME")
if server_name:
    app.config["SERVER_NAME"] = server_name
url_scheme = _get_env_str("MATRIX_URL_SCHEME")
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_port=1, x_prefix=1)
cache = Cache(app, config={"CACHE_TYPE": "SimpleCache"})
if _parse_bool_env("MATRIX_RATE_LIMIT_ENABLED", True):
    limiter = Limiter(get_remote_address, app=app, default_limits=[], storage_uri="memory://")
else:
    class UnlimitedLimiter:
        @staticmethod
        def limit(*args, **kwargs):
            return lambda f: f


    limiter = UnlimitedLimiter()
MAX_CONCURRENT_GENERATIONS = 2
MAX_QUEUE_SIZE = 5

gen_semaphore = Semaphore(MAX_CONCURRENT_GENERATIONS)
job_lock = Lock()


@dataclass(slots=True)
class GenerationJob:
    status: str
    filename: str
    session_id: str
    braille_preview: str
    created_at: float
    error: str | None = None


generation_jobs: dict[str, GenerationJob] = {}

MODELS_DIR = Path(__file__).parent / "models"
MODELS_DIR.mkdir(exist_ok=True)
TEMP_FILE_TTL_SECONDS = 600
GEN_TIMEOUT_SECONDS = 180
MAX_TEXT_LENGTH = 800


def _cleanup_expired_temporary_models() -> None:
    MODELS_DIR.mkdir(exist_ok=True)
    cutoff = time.time() - TEMP_FILE_TTL_SECONDS

    for entry in MODELS_DIR.iterdir():
        if not entry.is_file() or entry.suffix != ".stl":
            continue

        if entry.stat().st_mtime < cutoff:
            entry.unlink(missing_ok=True)

    with job_lock:
        expired_job_ids: list[str] = []
        for job_id, job in list(generation_jobs.items()):
            created_at = job.created_at
            status = job.status
            if isinstance(created_at, (int, float)) and created_at < cutoff and status != "running":
                expired_job_ids.append(job_id)

        for job_id in expired_job_ids:
            generation_jobs.pop(job_id, None)


def _get_or_create_session_id() -> str:
    existing_session_id = session.get("session_id")
    if isinstance(existing_session_id, str) and existing_session_id.strip():
        return existing_session_id

    new_session_id = uuid.uuid4().hex
    session["session_id"] = new_session_id
    return new_session_id


def _run_generation_job(
        job_id: str,
        text: str,
        output_path: Path,
        dot_radius: float | None,
        dot_spacing: float | None,
        row_spacing: float | None,
        column_spacing: float | None,
        page_thickness: float | None,
        max_page_width: float | None,
) -> None:
    gen_semaphore.acquire()

    with job_lock:
        job = generation_jobs.get(job_id)
        if job is None:
            gen_semaphore.release()
            return
        job.status = "running"

    try:
        generate_braille_model_from_text(
            text,
            str(output_path),
            dot_radius=dot_radius,
            dot_spacing=dot_spacing,
            row_spacing=row_spacing,
            column_spacing=column_spacing,
            page_thickness=page_thickness,
            max_page_width=max_page_width,
            gen_timeout_seconds=GEN_TIMEOUT_SECONDS,
        )
    except subprocess.TimeoutExpired:
        output_path.unlink(missing_ok=True)
        with job_lock:
            job = generation_jobs.get(job_id)
            if job is not None:
                job.status = "error"
                job.error = "Generation timed out."
    except Exception:
        output_path.unlink(missing_ok=True)
        with job_lock:
            job = generation_jobs.get(job_id)
            if job is not None:
                job.status = "error"
                job.error = "Generation failed."
    else:
        with job_lock:
            job = generation_jobs.get(job_id)
            if job is not None:
                job.status = "done"
                job.error = None
    finally:
        gen_semaphore.release()


@app.route("/livez", methods=["GET"])
def liveness():
    return jsonify({"status": "ok"}), 200


@app.route("/readyz", methods=["GET"])
def readiness():
    return jsonify({"status": "ready"}), 200


@app.route("/favicon.ico")
def favicon():
    return send_from_directory(app.static_folder, "favicon.ico")


@app.route("/")
@limiter.limit("300/minute", error_message="Please try again later.")
def index():
    _cleanup_expired_temporary_models()
    canonical_url = url_for("index", _scheme=url_scheme, _external=True)
    return render_template("index.html", canonical_url=canonical_url)


@app.route("/generate", methods=["POST"])
@limiter.limit("120/hour", error_message="Please try again later.")
@limiter.limit("20/hour;2/minute", key_func=_get_or_create_session_id, error_message="Please try again later.")
def generate():
    _cleanup_expired_temporary_models()
    text = request.form.get("text", "")
    if len(text) > MAX_TEXT_LENGTH:
        abort(400, description="Text exceeds maximum length.")
    braille_preview = translate_text_to_braille_text(text)
    session_id = _get_or_create_session_id()
    generation_id = uuid.uuid4().hex[:7]
    job_id = uuid.uuid4().hex
    output_filename = f"braille_model_{session_id[:7]}_{generation_id}.stl"
    output_path = MODELS_DIR / output_filename

    def get_float_or_none(key):
        value = request.form.get(key, "").strip()
        if not value:
            return None
        try:
            return float(value)
        except ValueError:
            abort(400, description="Invalid value.")

    dot_radius = get_float_or_none("dot_radius")
    dot_spacing = get_float_or_none("dot_spacing")
    row_spacing = get_float_or_none("row_spacing")
    column_spacing = get_float_or_none("column_spacing")
    page_thickness = get_float_or_none("page_thickness")
    max_page_width = get_float_or_none("max_page_width")

    with job_lock:
        pending_jobs = sum(1 for job in generation_jobs.values() if job.status == "pending")
        if pending_jobs >= MAX_QUEUE_SIZE:
            abort(503, description="Generation queue is full. Please try again later.")

        generation_jobs[job_id] = GenerationJob(
            status="pending",
            error=None,
            filename=output_filename,
            session_id=session_id,
            braille_preview=braille_preview,
            created_at=time.time(),
        )

    generation_thread = Thread(
        target=_run_generation_job,
        args=(
            job_id,
            text,
            output_path
        ),
        kwargs={
            "dot_radius": dot_radius,
            "dot_spacing": dot_spacing,
            "row_spacing": row_spacing,
            "column_spacing": column_spacing,
            "page_thickness": page_thickness,
            "max_page_width": max_page_width,
        },
        daemon=True
    )
    generation_thread.start()

    generated_files = session.get("generated_files")
    if not isinstance(generated_files, list):
        generated_files = []
    generated_files.append(output_filename)
    session["generated_files"] = generated_files[-3:]

    return render_template("download.html", filename=output_filename, job_id=job_id)


@app.route("/status/<job_id>")
@limiter.limit("420/minute", error_message="Please try again later.")
@limiter.limit("70/minute", key_func=_get_or_create_session_id, error_message="Please try again later.")
def generation_status(job_id):
    _cleanup_expired_temporary_models()
    session_id = _get_or_create_session_id()
    with job_lock:
        job = generation_jobs.get(job_id)

    if job is None or job.session_id != session_id:
        return jsonify({"status": "expired", "error": "Generation job not found."})

    status = job.status
    if not isinstance(status, str):
        status = "error"

    error = job.error
    if error is not None and not isinstance(error, str):
        error = "Generation failed."

    return jsonify(
        {
            "status": status,
            "error": error
        }
    )


@app.route("/preview/<job_id>")
@limiter.limit("420/minute", error_message="Please try again later.")
@limiter.limit("70/minute", key_func=_get_or_create_session_id, error_message="Please try again later.")
def generation_preview(job_id):
    _cleanup_expired_temporary_models()
    session_id = _get_or_create_session_id()
    with job_lock:
        job = generation_jobs.get(job_id)

    if job is None or job.session_id != session_id:
        return jsonify({"status": "expired", "error": "Generation job not found."})

    braille_preview = job.braille_preview
    if not isinstance(braille_preview, str):
        return jsonify({"status": "error", "error": "Braille preview unavailable."})

    return jsonify({"status": "ok", "braille_text": braille_preview})


@app.route("/download/<filename>")
@limiter.limit("120/minute", error_message="Please try again later.")
@limiter.limit("20/minute", key_func=_get_or_create_session_id, error_message="Please try again later.")
def download(filename):
    _cleanup_expired_temporary_models()
    generated_files = session.get("generated_files")
    if not isinstance(generated_files, list) or filename not in generated_files:
        abort(404, description="File not found or no longer available.")

    try:
        return send_from_directory(MODELS_DIR, filename, as_attachment=True)
    except FileNotFoundError:
        abort(404, description="File expired or already removed.")


@app.route("/robots.txt")
@cache.cached(timeout=0)
def robots():
    sitemap_url = url_for("sitemap", _scheme=url_scheme, _external=True)
    return Response(
        render_template("robots.txt", sitemap_url=sitemap_url),
        mimetype="text/plain"
    )


@app.route("/sitemap.xml")
@cache.cached(timeout=0)
def sitemap():
    pages = [
        url_for("index", _scheme=url_scheme, _external=True),
    ]
    return Response(
        render_template("sitemap.xml", pages=pages),
        mimetype="application/xml"
    )


if __name__ == "__main__":
    host = os.getenv("MATRIX_HOST", "0.0.0.0")
    port = _parse_int_env("MATRIX_PORT", 5000)
    debug = _parse_bool_env("MATRIX_DEBUG", False)
    app.run(host=host, port=port, debug=debug)