import argparse
import json
from pathlib import Path

from flask import Flask, Response, jsonify, request


app = Flask(__name__)
HITS = []
DEFAULT_FILENAME = "../../app/test.py"
DEFAULT_CONTENT = "print('stage1 demo hit')\n"
PROBE_MODEL = Path(__file__).resolve().parent / "trained_probe_model.tar.gz"


@app.get("/stage1")
def stage1():
    filename = request.args.get("filename", DEFAULT_FILENAME)
    content = request.args.get("content", DEFAULT_CONTENT)

    HITS.append(
        {
            "remote_addr": request.headers.get("X-Forwarded-For", request.remote_addr),
            "path": "/stage1",
            "filename": filename,
            "content_len": len(content),
        }
    )

    resp = Response(content.encode("utf-8"), mimetype="text/plain; charset=utf-8")
    resp.headers["filename"] = filename
    return resp


@app.get("/stage2-probe-model")
def stage2_probe_model():
    if not PROBE_MODEL.exists():
        return jsonify({"error": f"missing probe model: {PROBE_MODEL.name}"}), 500

    body = PROBE_MODEL.read_bytes()
    HITS.append(
        {
            "remote_addr": request.headers.get("X-Forwarded-For", request.remote_addr),
            "path": "/stage2-probe-model",
            "filename": PROBE_MODEL.name,
            "content_len": len(body),
        }
    )

    resp = Response(body, mimetype="application/x-tar")
    resp.headers["filename"] = PROBE_MODEL.name
    resp.headers["ETag"] = "trained-probe-model"
    return resp


@app.get("/hits")
def hits():
    return jsonify(HITS)


def main() -> None:
    parser = argparse.ArgumentParser(description="Minimal Flask server for filename write testing")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=5000)
    args = parser.parse_args()

    print(
        json.dumps(
            {
                "listen": f"http://{args.host}:{args.port}",
                "stage1_url": f"http://127.0.0.1:{args.port}/stage1",
                "stage2_probe_model_url": f"http://127.0.0.1:{args.port}/stage2-probe-model",
                "default_filename": DEFAULT_FILENAME,
            },
            ensure_ascii=False,
        ),
        flush=True,
    )

    app.run(host=args.host, port=args.port, debug=False)


if __name__ == "__main__":
    main()