#!/usr/bin/env python3
"""
Recounting upload dashboard — upload MP4 then launch recount via commander.
Consumes live counter + recount APIs, proxies recount live-video for preview.
"""
import json
import logging
import os
import re
import shlex
import signal
import subprocess
import threading
import time
import requests
from datetime import datetime
from pathlib import Path

from flask import Flask, render_template, jsonify, request, Response, send_file
from werkzeug.serving import WSGIRequestHandler
from werkzeug.utils import secure_filename

from dotenv import load_dotenv
load_dotenv()

log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

app = Flask(__name__, template_folder="templates")
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "change-me-in-production")
app.config["MAX_CONTENT_LENGTH"] = 512 * 1024 * 1024  # 512 MB

REPO_ROOT = Path(__file__).resolve().parent
UPLOAD_DIR = Path(os.getenv("UPLOAD_DIR", str(REPO_ROOT / "uploads")))
LIVE_API_URL = os.getenv("LIVE_API_URL", "http://localhost:5000")
RECOUNT_API_URL = os.getenv("RECOUNT_API_URL", "http://localhost:5001")

SITE_NAME = os.getenv("SITE_NAME", "RECOUNT")
TEMPLATE = os.getenv("RECOUNTING_UPLOAD_TEMPLATE", "recounting_upload_lamborghini.html")
LIVE_DASHBOARD_URL = os.getenv("LIVE_DASHBOARD_URL", "")

DASHBOARD_PORT = int(os.getenv("RECOUNTING_UPLOAD_PORT", "5003"))
DASHBOARD_HOST = os.getenv("DASHBOARD_HOST", "0.0.0.0")
FLASK_DEBUG = os.getenv("FLASK_DEBUG", "false").lower() == "true"

RECOUNT_CMD = os.getenv("RECOUNT_CMD", "bytetrack-counter config.env --source {path}")

SHM_DIR = os.getenv("SHM_DIR", "/dev/shm/bytetrack-counter")
_CONTINUE_FLAG = Path(f"{SHM_DIR}/.continue")

_http_session = requests.Session()
_http_session.timeout = 3

_FILENAME_RE = re.compile(r"batch[_-]?(\d+)[_-](\d{8})[_-]\d{6}", re.IGNORECASE)

_recount_proc = None
_recount_lock = threading.Lock()
_current_file = None


def _parse_filename(filename):
    m = _FILENAME_RE.search(filename)
    if m:
        raw = m.group(2)
        return int(m.group(1)), f"{raw[:4]}-{raw[4:6]}-{raw[6:8]}"
    return None, None


def _api_get(base_url, path, default=None):
    try:
        resp = _http_session.get(f"{base_url}{path}")
        if resp.status_code == 200:
            return resp.json()
    except Exception:
        pass
    return default


def _stop_recount_locked():
    global _recount_proc
    proc = _recount_proc
    _recount_proc = None
    if proc is None:
        return
    try:
        pgid = os.getpgid(proc.pid)
    except (ProcessLookupError, PermissionError):
        return
    try:
        os.killpg(pgid, signal.SIGTERM)
    except (ProcessLookupError, PermissionError):
        pass
    try:
        proc.wait(timeout=5)
    except subprocess.TimeoutExpired:
        try:
            os.killpg(pgid, signal.SIGKILL)
        except (ProcessLookupError, PermissionError):
            pass
        try:
            proc.wait(timeout=3)
        except Exception:
            pass
    except Exception:
        pass


def _folder_date(filepath, base):
    rel = Path(filepath).resolve().relative_to(base.resolve())
    parts = rel.parts
    if len(parts) >= 2:
        return parts[0]
    return None


@app.route("/")
def index():
    return render_template(TEMPLATE, site_name=SITE_NAME, upload_dir=str(UPLOAD_DIR), live_dashboard_url=LIVE_DASHBOARD_URL)


@app.route("/api/live-progress")
def api_live_progress():
    data = _api_get(LIVE_API_URL, "/api/current-batch")
    if data and data.get("success"):
        return jsonify(data)
    return jsonify({"success": False, "count": 0, "batch_number": None, "error": "Live unreachable"}), 200


@app.route("/api/recount-progress")
def api_recount_progress():
    data = _api_get(RECOUNT_API_URL, "/api/current-batch")
    if data and data.get("success"):
        return jsonify(data)
    return jsonify({"success": False, "count": 0, "batch_number": None, "error": "Recount unreachable"}), 200


@app.route("/api/uploads")
def api_uploads():
    files = []
    Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
    for f in sorted(Path(UPLOAD_DIR).glob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True):
        st = f.stat()
        files.append({
            "name": f.name,
            "path": str(f),
            "size": st.st_size,
            "mtime": datetime.fromtimestamp(st.st_mtime).isoformat(),
        })
    return jsonify(files)


@app.route("/api/upload", methods=["POST"])
def api_upload():
    if "file" not in request.files:
        return jsonify({"success": False, "error": "No file field"}), 400
    f = request.files["file"]
    if not f.filename:
        return jsonify({"success": False, "error": "No file selected"}), 400
    if not f.filename.lower().endswith(".mp4"):
        return jsonify({"success": False, "error": "Only .mp4 files allowed"}), 400

    Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
    dest_name = secure_filename(f.filename)
    dest_path = Path(UPLOAD_DIR) / dest_name

    f.save(str(dest_path))
    st = dest_path.stat()
    return jsonify({
        "success": True,
        "name": dest_name,
        "path": str(dest_path),
        "size": st.st_size,
    })


@app.route("/api/delete-upload", methods=["POST"])
def api_delete_upload():
    data = request.get_json(force=True) or {}
    fpath = data.get("path", "")
    if not fpath or not os.path.isfile(fpath):
        return jsonify({"success": False, "error": "File not found"}), 404
    f = Path(fpath)
    if f.parent.resolve() != Path(UPLOAD_DIR).resolve():
        return jsonify({"success": False, "error": "Not in upload dir"}), 403
    f.unlink()
    return jsonify({"success": True})


def _reset_recount_api():
    try:
        resp = _http_session.post(f"{RECOUNT_API_URL}/api/reset", timeout=5)
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Recount reset: {RECOUNT_API_URL}/api/reset → {resp.status_code}")
    except requests.exceptions.Timeout:
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Recount reset: {RECOUNT_API_URL}/api/reset → timeout")
    except requests.exceptions.ConnectionError as e:
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Recount reset: {RECOUNT_API_URL}/api/reset → unreachable: {e}")
    except Exception as e:
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Recount reset: {RECOUNT_API_URL}/api/reset → error: {e}")


def _spawn_recount_locked(upload_path):
    global _recount_proc
    cmd_str = RECOUNT_CMD.replace("{path}", shlex.quote(str(upload_path)))
    _recount_proc = subprocess.Popen(cmd_str, shell=True, start_new_session=True)
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Recount loaded: {cmd_str}")


@app.route("/api/load-mp4", methods=["POST"])
def load_mp4():
    global _current_file
    data = request.get_json(force=True) or {}
    upload_path = data.get("path", "")
    if not upload_path:
        return jsonify({"success": False, "error": "Missing 'path'"}), 400
    if not os.path.isfile(upload_path):
        return jsonify({"success": False, "error": f"File not found: {upload_path}"}), 404

    with _recount_lock:
        _stop_recount_locked()
        _current_file = upload_path
        try:
            _spawn_recount_locked(upload_path)
        except Exception as e:
            _current_file = None
            return jsonify({"success": False, "error": f"Failed to start: {e}"}), 500

    _reset_recount_api()

    return jsonify({
        "success": True,
        "stream_url": "/api/proxy-stream",
        "file": Path(upload_path).name,
        "path": upload_path,
    })


@app.route("/api/start-recount", methods=["POST"])
def start_recount():
    with _recount_lock:
        current = _current_file
    if not current:
        return jsonify({"success": False, "error": "Load an MP4 first"}), 400

    try:
        Path(SHM_DIR).mkdir(parents=True, exist_ok=True)
        _CONTINUE_FLAG.touch(exist_ok=True)
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Recount continue: {_CONTINUE_FLAG}")
    except Exception as e:
        return jsonify({"success": False, "error": f"Failed to signal: {e}"}), 500

    return jsonify({
        "success": True,
        "stream_url": "/api/proxy-stream",
        "file": Path(current).name,
        "path": current,
    })


@app.route("/api/stop-recount", methods=["POST"])
def stop_recount():
    global _current_file
    with _recount_lock:
        _stop_recount_locked()
        _current_file = None
    return jsonify({"success": True})


@app.route("/api/state")
def api_state():
    with _recount_lock:
        running = _recount_proc is not None and _recount_proc.poll() is None
        current = _current_file
    result = {
        "streaming": running,
    }
    if current:
        result["file"] = current
        result["file_name"] = Path(current).name
        result["stream_url"] = "/api/proxy-stream"
        batch_num, date_str = _parse_filename(Path(current).name)
        if batch_num is not None:
            result["batch_number"] = batch_num
            result["batch_date"] = date_str
    return jsonify(result)


@app.route("/api/proxy-stream")
def api_proxy_stream():
    try:
        resp = _http_session.get(f"{RECOUNT_API_URL}/api/live-video", stream=True, timeout=5)
        if resp.status_code != 200:
            return Response("stream unavailable", status=502)

        def generate():
            for chunk in resp.iter_content(chunk_size=8192):
                if chunk:
                    yield chunk
            resp.close()

        return Response(
            generate(),
            mimetype=resp.headers.get("Content-Type", "multipart/x-mixed-replace; boundary=frame"),
        )
    except Exception:
        return Response("stream unavailable", status=502)


@app.route("/api/batch-result")
def api_batch_result():
    fpath = request.args.get("file", "")
    filename = Path(fpath).name
    batch_num, date_str = _parse_filename(filename)
    if batch_num is None:
        return jsonify({"success": False, "error": "Cannot parse batch number from filename"}), 200
    if not date_str:
        date_str = _folder_date(fpath, Path(UPLOAD_DIR))
    if not date_str:
        return jsonify({"success": False, "error": "Cannot determine date from filename"}), 200

    def _match_in_batches(batches):
        for b in batches:
            if b.get("batch_number") == batch_num:
                return b
        return None

    data = _api_get(LIVE_API_URL, f"/api/day-detail/{date_str}")
    if data:
        b = _match_in_batches(data.get("batches", []))
        if b:
            return jsonify({
                "success": True,
                "batch_number": batch_num,
                "date": date_str,
                "count": b["count"],
                "start_time": b.get("start_time"),
                "end_time": b.get("end_time"),
            })

    recent = _api_get(LIVE_API_URL, "/api/recent-batches?limit=100")
    if recent:
        for entry in recent if isinstance(recent, list) else recent.get("batches", []):
            if isinstance(entry, dict) and entry.get("batch_number") == batch_num:
                return jsonify({
                    "success": True,
                    "batch_number": batch_num,
                    "date": entry.get("date", ""),
                    "count": entry["count"],
                    "start_time": entry.get("start_time"),
                    "end_time": entry.get("end_time"),
                })

    return jsonify({"success": False, "error": f"Batch #{batch_num} not found"}), 200


@app.route("/api/download-mp4")
def api_download_mp4():
    fpath = request.args.get("file", "")
    if not fpath or not os.path.isfile(fpath):
        return jsonify({"success": False, "error": "File not found"}), 404
    return send_file(fpath, as_attachment=True, download_name=Path(fpath).name)


if __name__ == "__main__":
    WSGIRequestHandler.protocol_version = "HTTP/1.1"
    Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
    print(f"Recounting upload dashboard at http://{DASHBOARD_HOST}:{DASHBOARD_PORT}")
    print(f"Live API: {LIVE_API_URL}")
    print(f"Recount API: {RECOUNT_API_URL}")
    print(f"Upload dir: {UPLOAD_DIR}")
    print(f"Recount command: {RECOUNT_CMD}")
    app.run(host=DASHBOARD_HOST, port=DASHBOARD_PORT, debug=FLASK_DEBUG)