# app.py
"""Flask web UI for feedmill_recounter — port 9000."""

from __future__ import annotations

import os

from dotenv import load_dotenv
from flask import (
    Flask, render_template, request, redirect,
    url_for, send_file, jsonify,
)
import time as _time
from flask import Response
from werkzeug.utils import secure_filename

from pathlib import Path

from src.job import JobQueue
from src.model_registry import scan_models, scan_model_groups, ModelConfig
from src.preview import extract_thumbnail, extract_sample_frames

load_dotenv()

app = Flask(__name__, template_folder="templates", static_folder="static")
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "change-me")
app.config["MAX_CONTENT_LENGTH"] = 2 * 1024 * 1024 * 1024  # 2GB

MODELS_DIR = os.getenv("MODELS_DIR", "./models")
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "./uploads")
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./output")

os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

job_queue = JobQueue(output_dir=OUTPUT_DIR)

VIDEO_EXTENSIONS = ('.mp4', '.avi', '.mkv', '.mov', '.webm')


def _parse_model_configs(form):
    """Parse model selection from request form.

    Returns:
        (model_configs, class_filters) tuple.
    """
    selected_stems = form.getlist("model_stems")
    selected_models = form.getlist("models")
    model_configs = []
    class_filters = {}

    if selected_stems:
        groups = scan_model_groups(MODELS_DIR)
        groups_by_stem = {g.stem: g for g in groups}

        for stem in selected_stems:
            if stem not in groups_by_stem:
                continue
            group = groups_by_stem[stem]
            fmt = form.get(f"format_{stem}", group.default_format)
            if fmt not in group.format_paths:
                continue
            model_configs.append(ModelConfig(
                filename=os.path.basename(group.format_paths[fmt]),
                path=group.format_paths[fmt],
                stem=stem,
                known_classes=list(group.known_classes),
            ))
            filter_val = form.get(f"filter_{stem}", "")
            if not filter_val or filter_val in ("default",):
                pass
            elif filter_val == "all":
                class_filters[stem] = None
            else:
                class_filters[stem] = filter_val.split(",")
    elif selected_models:
        models = scan_models(MODELS_DIR)
        by_name = {m.filename: m for m in models}

        for name in selected_models:
            if name in by_name:
                model_configs.append(by_name[name])
                filter_val = form.get(f"filter_{name}", "")
                if not filter_val or filter_val in ("default",):
                    pass
                elif filter_val == "all":
                    class_filters[name] = None
                else:
                    class_filters[name] = filter_val.split(",")

    return model_configs, class_filters


@app.template_filter("basename")
def basename_filter(path):
    """Extract filename from path for templates."""
    return os.path.basename(path)


@app.route("/")
def index():
    model_groups = scan_model_groups(MODELS_DIR)
    return render_template("index.html", model_groups=model_groups, models_dir=MODELS_DIR)


@app.route("/upload", methods=["POST"])
def upload():
    video = request.files.get("video")
    if not video or not video.filename:
        return "No video uploaded", 400

    safe_name = secure_filename(video.filename)
    if not safe_name or not safe_name.lower().endswith(VIDEO_EXTENSIONS):
        return "Invalid video file type", 400

    model_configs, class_filters = _parse_model_configs(request.form)
    if not model_configs:
        return "No models selected", 400

    video_path = os.path.join(UPLOAD_DIR, safe_name)
    base, ext = os.path.splitext(video_path)
    n = 1
    while os.path.exists(video_path):
        video_path = f"{base}_{n}{ext}"
        n += 1
    video.save(video_path)

    job = job_queue.add_job(
        video_path=video_path,
        model_configs=model_configs,
        class_filters=class_filters,
    )
    return redirect(url_for("status", job_id=job.job_id))


@app.route("/upload/reuse", methods=["POST"])
def upload_reuse():
    """Re-analyze an existing uploaded video."""
    video_path = request.form.get("video_path", "")
    if not video_path or not os.path.isfile(video_path):
        return "Video not found", 400

    model_configs, class_filters = _parse_model_configs(request.form)
    if not model_configs:
        return "No models selected", 400

    job = job_queue.add_job(
        video_path=video_path,
        model_configs=model_configs,
        class_filters=class_filters,
    )
    return redirect(url_for("status", job_id=job.job_id))


@app.route("/status/<job_id>")
def status(job_id):
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    return render_template("status.html", job=job)


@app.route("/jobs")
def jobs_list():
    jobs = job_queue.list_jobs()
    return render_template("jobs.html", jobs=jobs)


@app.route("/download/<job_id>/<filename>")
def download(job_id, filename):
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    safe_filename = secure_filename(filename)
    output_dir_abs = os.path.abspath(job.output_dir)
    file_path = os.path.abspath(os.path.join(job.output_dir, safe_filename))
    if os.path.commonpath([output_dir_abs, file_path]) != output_dir_abs:
        return "File not found", 404
    if not os.path.isfile(file_path):
        return "File not found", 404
    return send_file(file_path, as_attachment=True)


@app.route("/api/videos")
def api_videos():
    """List previously uploaded videos."""
    videos = []
    for f in sorted(Path(UPLOAD_DIR).iterdir(), key=lambda x: x.stat().st_mtime, reverse=True):
        if f.is_file() and f.suffix.lower() in VIDEO_EXTENSIONS:
            videos.append({
                "filename": f.name,
                "size": f.stat().st_size,
                "mtime": f.stat().st_mtime,
                "path": str(f),
            })
    return jsonify(videos)


@app.route("/api/videos/<filename>", methods=["DELETE"])
def api_delete_video(filename):
    """Delete an uploaded video."""
    safe = secure_filename(filename)
    if not safe or not safe.lower().endswith(VIDEO_EXTENSIONS):
        return jsonify({"error": "invalid filename"}), 400
    file_path = os.path.join(UPLOAD_DIR, safe)
    if not os.path.isfile(file_path):
        return jsonify({"error": "not found"}), 404
    os.remove(file_path)
    return jsonify({"status": "deleted"})


@app.route("/api/models")
def api_models():
    groups = scan_model_groups(MODELS_DIR)
    return jsonify([
        {
            "stem": g.stem,
            "formats": g.formats,
            "format_paths": g.format_paths,
            "known_classes": g.known_classes,
            "default_format": g.default_format,
        }
        for g in groups
    ])


@app.route("/api/jobs")
def api_jobs():
    return jsonify([{
        "job_id": j.job_id,
        "status": j.status.name,
        "progress": j.progress,
        "video_path": os.path.basename(j.video_path),
        "results": [
            {
                "model": r.model_name,
                "loading": r.loading_count,
                "unloading": r.unloading_count,
                "net": r.net_count,
            }
            for r in j.results
        ],
    } for j in job_queue.list_jobs()])


@app.route("/api/jobs/<job_id>/cancel", methods=["POST"])
def api_cancel_job(job_id):
    """Cancel a running or pending job."""
    success = job_queue.cancel_job(job_id)
    if success:
        return jsonify({"status": "cancelled"})
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    return jsonify({"error": "cannot cancel job in state " + job.status.name}), 400


@app.route("/api/jobs/<job_id>")
def api_job_detail(job_id):
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    return jsonify({
        "job_id": job.job_id,
        "status": job.status.name,
        "progress": job.progress,
        "current_model": job.current_model,
        "results": [
            {
                "model_name": r.model_name,
                "loading_count": r.loading_count,
                "unloading_count": r.unloading_count,
                "net_count": r.net_count,
                "batch_count": r.batch_count,
                "frame_count": r.frame_count,
                "duration_seconds": r.duration_seconds,
                "output_path": os.path.basename(r.output_path) if r.output_path else None,
            }
            for r in job.results
        ],
        "error": job.error,
    })


@app.route("/preview/thumb/<job_id>")
def preview_thumb(job_id):
    """Serve a thumbnail JPEG for a job's video (cached in job output_dir)."""
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    thumb_path = os.path.join(job.output_dir, "thumbnail.jpg")
    if not os.path.isfile(thumb_path):
        try:
            extract_thumbnail(job.video_path, thumb_path)
        except (FileNotFoundError, RuntimeError):
            return "Could not generate thumbnail", 404
    return send_file(thumb_path, mimetype="image/jpeg")


@app.route("/api/jobs/<job_id>/samples")
def api_job_samples(job_id):
    """Return JSON array of sample frame paths (extract if not cached)."""
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    samples_dir = os.path.join(job.output_dir, "samples")
    cache_file = os.path.join(samples_dir, ".cached")
    if not os.path.isfile(cache_file):
        try:
            paths = extract_sample_frames(job.video_path, samples_dir)
            with open(cache_file, "w") as f:
                f.write("\n".join(paths))
        except (FileNotFoundError, RuntimeError):
            return jsonify({"error": "could not generate samples"}), 404
    else:
        with open(cache_file) as f:
            paths = [line.strip() for line in f if line.strip()]
    rel_paths = [os.path.relpath(p, job.output_dir) for p in paths]
    return jsonify(rel_paths)


@app.route("/api/preview/<job_id>")
def api_preview(job_id):
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    preview_path = job.preview_path
    if not preview_path:
        return jsonify({"error": "no preview path"}), 404

    def generate():
        consecutive_fails = 0
        MAX_FAILS = 30
        while True:
            try:
                with open(preview_path, "rb") as f:
                    jpeg = f.read()
                consecutive_fails = 0
                yield (b"--frame\r\n"
                       b"Content-Type: image/jpeg\r\n\r\n" + jpeg + b"\r\n")
            except FileNotFoundError:
                consecutive_fails += 1
                if consecutive_fails >= MAX_FAILS:
                    return
                _time.sleep(1.0)
                continue
            except Exception:
                consecutive_fails += 1
                if consecutive_fails >= MAX_FAILS:
                    return
                _time.sleep(0.5)
                continue
            _time.sleep(0.05)
    return Response(generate(), mimetype="multipart/x-mixed-replace; boundary=frame")


def main():
    host = os.getenv("WEB_HOST", "0.0.0.0")
    port = int(os.getenv("WEB_PORT", "9000"))
    debug = os.getenv("FLASK_DEBUG", "false").lower() == "true"

    print(f"Feedmill Recounter web UI: http://{host}:{port}")
    app.run(host=host, port=port, debug=debug)


if __name__ == "__main__":
    main()