import io
import os
import random
import time
import base64
from pathlib import Path

import pandas as pd
import torch
import torchvision.transforms as T
from flask import Flask, jsonify, render_template, request, session
from PIL import Image

app = Flask(__name__, template_folder="templates")
app.secret_key = os.environ.get("SECRET_KEY", "horse3903-dev-key")

FLAG = os.environ.get("FLAG", "sctf{REDACTED}")
BASE = Path(__file__).parent
MODELS_DIR = BASE / "models"

CLASSES = ["plane", "car", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
HORSE_IDX = 7
TARGET_CLASS = "ship"
TARGET_IDX = 8
SESSION_TIMEOUT = 120
MAX_LEVELS = 10
MAX_CONSECUTIVE_FAILURES = # redacted
SIMILARITY_THRESHOLD = # redacted

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

NORM_MEAN = # redacted
NORM_STD = # redacted

_to_tensor = T.ToTensor()
_norm = T.Normalize(NORM_MEAN, NORM_STD)

_lpips_fn = # redacted


def _build_residual_network():
    # redacted
    return


def _build_verifier():
    # redacted
    return


def _load_models():
    _m1 = _build_residual_network()
    _m1.load_state_dict(torch.load(MODELS_DIR / "bouncer.pt", map_location="cpu"))
    _m1.eval().to(DEVICE)

    _m2 = _build_verifier()
    _m2.load_state_dict(torch.load(MODELS_DIR / "model_2.pt", map_location="cpu"))
    _m2.eval().to(DEVICE)

    return _m1, _m2


def _load_horse_images():
    df = pd.read_parquet(BASE / "data" / "test.parquet")
    return [
        Image.open(io.BytesIO(row["img"]["bytes"])).convert("RGB")
        for _, row in df[df["label"] == HORSE_IDX].iterrows()
    ]


print("Loading models...", flush=True)
_m1, _m2 = _load_models()
print("Loading horse images...", flush=True)
horse_images = _load_horse_images()
print(f"Ready. {len(horse_images)} horse images loaded.", flush=True)


def _pil_to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def _preprocess(img: Image.Image) -> torch.Tensor:
    # redacted
    ...


def _preprocess_similarity(img: Image.Image) -> torch.Tensor:
    # redacted
    ...


def _classify(model, tensor: torch.Tensor) -> int:
    with torch.no_grad():
        return int(model(tensor).argmax(dim=1).item())


def _check_submission(original: Image.Image, submitted: Image.Image) -> str | None:
    """Run all checks. Returns an error key on failure, None on pass."""
    # redacted
    ...


def _new_round() -> tuple[int, str]:
    idx = random.randrange(len(horse_images))
    return idx, TARGET_CLASS


def _fail_response(msg: str, kicked: bool = False) -> dict:
    if kicked:
        idx, target = _new_round()
        session["level"] = 0
        session["consecutive_failures"] = 0
        session["img_idx"] = idx
        session["target_class"] = target
        img_b64 = _pil_to_b64(horse_images[idx])
        return {"status": "fail", "message": msg, "level": 0, "target": target, "img_b64": img_b64}
    return {
        "status": "fail",
        "message": msg,
        "level": session.get("level", 0),
        "target": TARGET_CLASS,
        "img_b64": None,
    }


def _bump_failures() -> bool:
    session["consecutive_failures"] = session.get("consecutive_failures", 0) + 1
    return session["consecutive_failures"] > MAX_CONSECUTIVE_FAILURES


@app.route("/")
def index():
    now = time.time()

    if "start_time" in session and now - session["start_time"] > SESSION_TIMEOUT:
        session.clear()

    if "level" not in session:
        session["level"] = 0
        session["consecutive_failures"] = 0
        session["solved"] = False

    if "img_idx" not in session:
        idx, target = _new_round()
        session["img_idx"] = idx
        session["target_class"] = target
        session["start_time"] = now

    img = horse_images[session["img_idx"]]
    remaining = max(0, int(SESSION_TIMEOUT - (now - session.get("start_time", now))))

    return render_template(
        "index.html",
        img_b64=_pil_to_b64(img),
        target=session["target_class"],
        level=session["level"],
        max_levels=MAX_LEVELS,
        remaining=remaining,
    )


@app.route("/submit", methods=["POST"])
def submit():
    now = time.time()

    if "start_time" not in session:
        return jsonify({"status": "fail", "message": "No active operation. Refresh to begin.", "level": 0, "target": "", "img_b64": None})

    if now - session["start_time"] > SESSION_TIMEOUT:
        session.clear()
        return jsonify({"status": "fail", "message": "The gala has begun. Your window has closed.", "level": 0, "target": "", "img_b64": None})

    if "file" not in request.files:
        return jsonify(_fail_response("No forgery submitted."))

    try:
        submitted = Image.open(request.files["file"]).convert("RGB")
    except Exception:
        return jsonify(_fail_response("The submission could not be read as an image."))

    original = horse_images[session["img_idx"]]
    target_class = TARGET_CLASS

    err = _check_submission(original, submitted)

    if err == "similarity":
        kicked = _bump_failures()
        return jsonify(_fail_response(
            "Security has been alerted. Escorted from the premises." if kicked
            else "The conservator has pulled this work. Alterations are visible.",
            kicked=kicked,
        ))

    if err == "verifier":
        kicked = _bump_failures()
        return jsonify(_fail_response(
            "Security has been alerted. Escorted from the premises." if kicked
            else "The conservator no longer sees an equestrian subject. This is no longer a convincing forgery.",
            kicked=kicked,
        ))

    if err == "bouncer_horse":
        kicked = _bump_failures()
        return jsonify(_fail_response(
            "Security has been alerted. Escorted from the premises." if kicked
            else "The cataloguing system flagged this as equestrian. Forgery detected.",
            kicked=kicked,
        ))

    if err == "bouncer_wrong":
        kicked = _bump_failures()
        return jsonify(_fail_response(
            "Security has been alerted. Escorted from the premises." if kicked
            else f"Catalogued as '{CLASSES[_classify(_m1, _preprocess(submitted))]}', not '{target_class}'. The client's contact is in the wrong wing.",
            kicked=kicked,
        ))

    session["consecutive_failures"] = 0
    session["level"] = session.get("level", 0) + 1
    new_level = session["level"]

    if new_level >= MAX_LEVELS:
        session["solved"] = True
        return jsonify({"status": "flag", "message": f"All {MAX_LEVELS} works have been catalogued and hung. The gala is underway. In the confusion, the switch is made.\n{FLAG}", "level": new_level, "target": target_class, "img_b64": None})

    idx, next_target = _new_round()
    session["img_idx"] = idx
    session["target_class"] = next_target
    img_b64 = _pil_to_b64(horse_images[idx])

    return jsonify({
        "status": "success",
        "message": f"Filed under: {target_class}. Forgery accepted. ({new_level}/{MAX_LEVELS})",
        "level": new_level,
        "target": next_target,
        "img_b64": img_b64,
    })


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=False)
