import os
from pathlib import Path

from flask import Flask, request

app = Flask(__name__)

_mode = "normal"
_first_fail_remaining = False
_uploads: dict[tuple[str, str], bytes] = {}
_load_count = 0
_upload_count = 0


def _models_root() -> Path:
    return Path(os.environ.get("MODELS_ROOT", "/models"))


def _resolve_disk_path(filename: str, folder: str | None) -> Path | None:
    root = _models_root()
    if folder:
        p = root / folder / filename
    else:
        p = root / filename
    if p.is_file():
        return p
    if folder is None:
        alt = root / "models" / filename
        if alt.is_file():
            return alt
    return None


def _should_fail_load() -> bool:
    global _first_fail_remaining
    if _mode == "error":
        return True
    if _mode == "first_fail":
        if _first_fail_remaining:
            _first_fail_remaining = False
            return True
        return False
    return False


@app.route("/load/<path:filename>", methods=["GET", "POST"])
def load(filename):
    global _load_count
    folder = None
    if request.method == "POST" and request.is_json:
        body = request.get_json(silent=True) or {}
        folder = body.get("folder")
    if _should_fail_load():
        return "", 503
    path = _resolve_disk_path(filename, folder)
    if path is None:
        key = (folder or "", filename)
        data = _uploads.get(key)
        if data is None and folder:
            data = _uploads.get(("", filename))
        if data is None:
            return "", 404
        _load_count += 1
        return data, 200
    _load_count += 1
    return path.read_bytes(), 200


@app.route("/upload/<path:filename>", methods=["POST"])
def upload(filename):
    global _upload_count
    folder = request.form.get("folder") or ""
    f = request.files.get("data")
    if not f:
        return "", 400
    _uploads[(folder, filename)] = f.read()
    _upload_count += 1
    return "", 200


@app.route("/mock/config", methods=["POST"])
def mock_config():
    global _mode, _first_fail_remaining
    body = request.get_json(silent=True) or {}
    mode = body.get("mode", "normal")
    if mode not in ("normal", "error", "first_fail"):
        return "", 400
    _mode = mode
    _first_fail_remaining = mode == "first_fail"
    return "", 200


@app.route("/mock/reset", methods=["POST"])
def mock_reset():
    global _mode, _first_fail_remaining, _uploads, _load_count, _upload_count
    _mode = "normal"
    _first_fail_remaining = False
    _uploads.clear()
    _load_count = 0
    _upload_count = 0
    return "", 200


@app.route("/mock/status", methods=["GET"])
def mock_status():
    return {
        "mode": _mode,
        "upload_count": _upload_count,
        "load_count": _load_count,
    }