"""BG Remover Web - Flask app (runs on cPanel "Setup Python App" or anywhere with WSGI)."""
import io
import json
import os
import re
import shutil
import time
import uuid
import zipfile

from flask import Flask, g, jsonify, request, send_file, send_from_directory
from PIL import Image, ImageOps

import config
import engine
import licensing

Image.MAX_IMAGE_PIXELS = config.MAX_PIXELS * 2
app = Flask(__name__, static_folder=None)
app.config["MAX_CONTENT_LENGTH"] = config.MAX_UPLOAD_MB * 1024 * 1024
os.makedirs(config.JOBS_DIR, exist_ok=True)

STATIC = os.path.join(config.BASE, "static")
EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".jfif")
JOB_RE = re.compile(r"^[0-9a-f]{32}$")
MID_RE = re.compile(r"^web-[0-9a-f]{32}$")


# ---------- helpers ----------
def client_ip():
    xff = request.headers.get("X-Forwarded-For", "").split(",")[0].strip()
    return xff or request.remote_addr or ""


def device():
    """Id of this browser (cookie). It plays the role of the PC's machine id."""
    mid = request.cookies.get("bgmid", "")
    if MID_RE.match(mid):
        return mid
    if not hasattr(g, "new_mid"):
        g.new_mid = "web-" + uuid.uuid4().hex
    return g.new_mid


@app.after_request
def after(resp):
    if hasattr(g, "new_mid"):
        secure = request.is_secure or request.headers.get("X-Forwarded-Proto", "") == "https"
        resp.set_cookie("bgmid", g.new_mid, max_age=3 * 365 * 86400, httponly=True, samesite="Lax", secure=secure)
    if request.path.startswith("/api/"):
        resp.headers["Cache-Control"] = "no-store"
    return resp


@app.errorhandler(413)
def too_big(_e):
    return jsonify(ok=False, error="toobig"), 413


def err(code, status=400, **kw):
    return jsonify(dict(ok=False, error=code, **kw)), status


def job_dir(job, must_exist=True):
    if not JOB_RE.match(str(job or "")):
        return None
    d = os.path.join(config.JOBS_DIR, job)
    if must_exist and not os.path.isdir(d):
        return None
    return d


def cleanup_jobs():
    now = time.time()
    try:
        for n in os.listdir(config.JOBS_DIR):
            d = os.path.join(config.JOBS_DIR, n)
            if now - os.path.getmtime(d) > config.JOB_TTL:
                shutil.rmtree(d, ignore_errors=True)
    except Exception:
        pass


def parse_bg(s):
    s = str(s or "").strip().lstrip("#")
    if re.fullmatch(r"[0-9a-fA-F]{6}", s):
        return tuple(int(s[i:i + 2], 16) for i in (0, 2, 4))
    return None


def fnum(v, default, lo, hi):
    try:
        return max(lo, min(hi, float(v)))
    except Exception:
        return default


def flag(v):
    return str(v).lower() in ("1", "true", "yes", "on")


def allowed(mid, touch=False):
    st = licensing.web_status(mid, touch=touch)
    return st["mode"] != "none", st


# ---------- pages ----------
@app.route("/")
def index():
    return send_from_directory(STATIC, "index.html")


@app.route("/static/<path:name>")
def static_files(name):
    return send_from_directory(STATIC, name)


@app.route("/healthz")
def healthz():
    return "ok"


# ---------- license API for the desktop app + keygen.py (same protocol as the old Google script) ----------
@app.route("/license", methods=["GET", "POST"])
def license_api():
    out = licensing.route(request.values.to_dict(), client_ip())
    return jsonify(out)


# ---------- website API ----------
@app.route("/api/state")
def api_state():
    mid = device()
    st = licensing.web_status(mid)
    st.update(ok=True, buy_url=licensing.buy_url(), models=engine.available_models(),
              default_model=engine.MODEL_NAME, trial_hours=config.TRIAL_HOURS)
    return jsonify(st)


@app.route("/api/activate", methods=["POST"])
def api_activate():
    d = request.get_json(silent=True) or {}
    mid = device()
    res = licensing.run("activate", {"code": d.get("code"), "name": str(d.get("name") or "").strip(), "mid": mid,
                                      "pc": "Web " + client_ip()}, client_ip())
    return jsonify(res)


@app.route("/api/trial", methods=["POST"])
def api_trial():
    mid = device()
    return jsonify(licensing.run("trial", {"mid": mid, "pc": "Web " + client_ip()}, client_ip()))


@app.route("/api/deactivate", methods=["POST"])
def api_deactivate():
    mid = device()
    for code in licensing.web_code_for(mid):
        licensing.run("deactivate", {"code": code, "mid": mid})
    return jsonify(ok=True)


@app.route("/api/process", methods=["POST"])
def api_process():
    mid = device()
    ok, st = allowed(mid, touch=True)
    if not ok:
        return err("license", 403, reason=st.get("reason", ""))
    cleanup_jobs()

    f = request.form
    mode = f.get("mode", "auto")
    if mode not in ("auto", "ai", "solid", "ink"):
        mode = "auto"
    tol = fnum(f.get("tol"), 30, 5, 90)
    shrink = int(round(fnum(f.get("shrink"), 1, 0, 6)))
    smooth = float(round(fnum(f.get("smooth"), 1, 0, 4)))
    models = engine.available_models()
    model = f.get("model") if f.get("model") in models else models[0]
    tol_manual, model_manual = flag(f.get("tol_manual")), flag(f.get("model_manual"))

    job = f.get("job", "")
    upload = request.files.get("file")
    try:
        if upload is not None:
            if not (upload.filename or "").lower().endswith(EXTS):
                return err("format")
            img = Image.open(upload.stream)
            if img.width * img.height > config.MAX_PIXELS:
                return err("toobig")
            img = ImageOps.exif_transpose(img).convert("RGBA")
            if max(img.size) > config.MAX_SIDE:
                img.thumbnail((config.MAX_SIDE, config.MAX_SIDE), Image.LANCZOS)
            job = uuid.uuid4().hex
            d = job_dir(job, must_exist=False)
            os.makedirs(d, exist_ok=True)
            img.save(os.path.join(d, "src.png"))
        else:
            d = job_dir(job)
            if not d:
                return err("job")
            img = Image.open(os.path.join(d, "src.png")).convert("RGBA")
    except Exception:
        return err("format")

    try:
        with engine.Busy():
            out, raw, picked = engine.run(img, mode, tol, model, shrink, smooth, tol_manual, model_manual)
    except FileNotFoundError as e:
        return err("model", 500, message=str(e))
    except MemoryError:
        return err("memory", 500)
    except Exception as e:                                   # noqa: BLE001
        return err("server", 500, message=str(e)[-300:])

    out.save(os.path.join(d, "result.png"))
    rawp = os.path.join(d, "raw.png")
    if raw is not None:
        raw.save(rawp)
    elif os.path.exists(rawp):
        os.remove(rawp)
    os.utime(d, None)
    return jsonify(ok=True, job=job, picked=picked, w=out.width, h=out.height, has_raw=raw is not None)


@app.route("/api/refine", methods=["POST"])
def api_refine():
    mid = device()
    ok, st = allowed(mid)
    if not ok:
        return err("license", 403, reason=st.get("reason", ""))
    d = request.get_json(silent=True) or {}
    jd = job_dir(d.get("job"))
    if not jd:
        return err("job")
    rawp = os.path.join(jd, "raw.png")
    if not os.path.exists(rawp):                              # solid / ink result: nothing to refine
        return jsonify(ok=True, changed=False)
    shrink = int(round(fnum(d.get("shrink"), 1, 0, 6)))
    smooth = float(round(fnum(d.get("smooth"), 1, 0, 4)))
    with engine.Busy():
        out = engine.refine_edges(Image.open(rawp).convert("RGBA"), shrink, smooth)
    out.save(os.path.join(jd, "result.png"))
    os.utime(jd, None)
    return jsonify(ok=True, changed=True)


@app.route("/api/img/<job>/<which>")
def api_img(job, which):
    jd = job_dir(job)
    if not jd or which not in ("src", "result"):
        return err("job", 404)
    p = os.path.join(jd, which + ".png")
    if not os.path.exists(p):
        return err("job", 404)
    if request.args.get("thumb"):
        im = Image.open(p)
        im.thumbnail((1400, 1400))
        buf = io.BytesIO()
        im.save(buf, "PNG")
        buf.seek(0)
        return send_file(buf, mimetype="image/png")
    return send_file(p, mimetype="image/png")


def render(result_path, fmt, bg, inch):
    """Result image bytes in the requested form (same rules as Save PNG / Save JPEG / Copy for Word)."""
    res = Image.open(result_path).convert("RGBA")
    dpi = engine.dpi_for(res.width, inch)
    buf = io.BytesIO()
    if fmt == "jpg":
        engine.compose_bg(res, bg or (255, 255, 255)).save(buf, "JPEG", quality=95, subsampling=0, dpi=dpi)
    else:
        (res if bg is None else engine.compose_bg(res, bg)).save(buf, "PNG", dpi=dpi)
    return buf.getvalue()


def safe_name(n):
    n = re.sub(r"[^\w\-. ]+", "_", str(n or "image"), flags=re.UNICODE).strip(" .") or "image"
    return n[:80]


@app.route("/api/download/<job>")
def api_download(job):
    mid = device()
    ok, st = allowed(mid)
    if not ok:
        return err("license", 403)
    jd = job_dir(job)
    if not jd:
        return err("job", 404)
    fmt = "jpg" if request.args.get("fmt") == "jpg" else "png"
    data = render(os.path.join(jd, "result.png"), fmt, parse_bg(request.args.get("bg")), request.args.get("inch"))
    name = safe_name(request.args.get("name")) + ("_nobg.jpg" if fmt == "jpg" else "_nobg.png")
    return send_file(io.BytesIO(data), mimetype="image/jpeg" if fmt == "jpg" else "image/png",
                     as_attachment=request.args.get("inline") != "1", download_name=name)


@app.route("/api/zip", methods=["POST"])
def api_zip():
    mid = device()
    ok, st = allowed(mid)
    if not ok:
        return err("license", 403)
    d = request.get_json(silent=True) or {}
    bg, inch = parse_bg(d.get("bg")), d.get("inch")
    buf = io.BytesIO()
    used = set()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as z:
        for it in (d.get("items") or [])[:200]:
            jd = job_dir(it.get("job"))
            if not jd:
                continue
            base = safe_name(it.get("name"))
            n, k = base, 2
            while n in used:
                n, k = "%s_%d" % (base, k), k + 1
            used.add(n)
            z.writestr(n + "_nobg.png", render(os.path.join(jd, "result.png"), "png", bg, inch))
    buf.seek(0)
    return send_file(buf, mimetype="application/zip", as_attachment=True, download_name="bg_removed.zip")


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=int(os.environ.get("PORT", "8000")), debug=False)