#!/usr/bin/env python3
"""
Edge Jetson production counter dashboard.
Reads jetson_counter.db + current_counter.json from the counter stack.
Tracks daily counter_in / counter_out and total per counting day (no batches).
Default port 5000.
"""

import json
import os
import re
import sqlite3
import time
from io import BytesIO
from pathlib import Path
from datetime import datetime, timedelta

from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, PatternFill, Border, Side

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

from dotenv import load_dotenv
load_dotenv()


app = Flask(__name__, template_folder="templates")
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "change-me-in-production")

_DEFAULT_DIR = "/opt/jetson-counter"
DB_PATH = os.getenv("DB_PATH", f"{_DEFAULT_DIR}/jetson_counter.db")
CURRENT_COUNTER_PATH = os.getenv("STATE_FILE", os.getenv("CURRENT_COUNTER_PATH", f"{_DEFAULT_DIR}/current_counter.json"))
CUTOFF_TIME = os.getenv("CUTOFF_TIME", os.getenv("DAILY_CUTOFF_TIME", "20:00"))

LIVE_STREAM_FRAME_PATH = os.getenv("LIVE_STREAM_FRAME_PATH", "/dev/shm/jetson-counter/live_frame.jpg")

CROSS_SNAPSHOT_DIR = os.getenv("CROSS_SNAPSHOT_DIR", f"{_DEFAULT_DIR}/snapshots")
SAVE_DETECT_SNAPSHOT = os.getenv("SAVE_DETECT_SNAPSHOT", "false").lower() == "true"

CONTROL_ENABLED = os.getenv("CONTROL_ENABLED", "false").lower() == "true"
CONTROL_FILE = os.getenv("CONTROL_FILE", f"{_DEFAULT_DIR}/control.json")
CONTROL_DEFAULT_COUNTING = os.getenv("CONTROL_DEFAULT_COUNTING", "true").lower() == "true"

SITE_NAME = os.getenv("SITE_NAME", "LIVE")

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

_JPEG_SOI = b"\xff\xd8"
_JPEG_EOI = b"\xff\xd9"
_MIN_JPEG_BYTES = 128


def _is_valid_jpeg(data):
    return (
        data
        and len(data) >= _MIN_JPEG_BYTES
        and data.startswith(_JPEG_SOI)
        and data.endswith(_JPEG_EOI)
    )


@app.route("/api/live-video")
def api_live_video():
    def generate():
        while True:
            try:
                if not os.path.isfile(LIVE_STREAM_FRAME_PATH):
                    time.sleep(0.5)
                    continue
                with open(LIVE_STREAM_FRAME_PATH, "rb") as f:
                    jpeg = f.read()
                if not jpeg or len(jpeg) < 2 or jpeg[:2] != b"\xff\xd8":
                    consecutive_fails += 1
                    if consecutive_fails >= MAX_FAILS:
                        return
                    time.sleep(0.05)
                    continue
                consecutive_fails = 0
                yield (b"--frame\r\n"
                       b"Content-Type: image/jpeg\r\n\r\n" + jpeg + b"\r\n")
            except FileNotFoundError:
                time.sleep(0.5)
                continue
            except Exception:
                time.sleep(0.25)
                continue
            time.sleep(0.05)
    return Response(generate(), mimetype="multipart/x-mixed-replace; boundary=frame")


_SNAP_RE = re.compile(
    r"^(?P<ts>\d{8}_\d{6}_\d{3})_(?P<kind>detect|in|out)_id(?P<tid>\d+)_f(?P<frame>\d+)\.jpg$"
)


def _parse_snapshot(path, category):
    m = _SNAP_RE.match(path.name)
    if not m:
        return None
    try:
        dt = datetime.strptime(m.group("ts"), "%Y%m%d_%H%M%S_%f")
    except ValueError:
        dt = datetime.fromtimestamp(path.stat().st_mtime)
    kind = m.group("kind")
    return {
        "file": f"{category}/{path.name}",
        "category": category,
        "kind": kind,
        "track_id": int(m.group("tid")),
        "frame": int(m.group("frame")),
        "timestamp": dt.isoformat(),
        "mtime": path.stat().st_mtime,
    }


def _collect_snapshots():
    base = os.path.abspath(CROSS_SNAPSHOT_DIR)
    items = []
    for category in ("cross", "detect"):
        sub = os.path.join(base, category)
        if not os.path.isdir(sub):
            continue
        for name in os.listdir(sub):
            if not name.lower().endswith(".jpg"):
                continue
            info = _parse_snapshot(Path(sub) / name, category)
            if info:
                items.append(info)
    items.sort(key=lambda x: x["mtime"], reverse=True)
    return items


@app.route("/api/snapshots")
def api_snapshots():
    try:
        kind = request.args.get("kind", "all")
        track_id = request.args.get("track_id", type=int)
        date = request.args.get("date")
        limit = request.args.get("limit", 200, type=int)

        items = _collect_snapshots()
        if kind and kind != "all":
            if kind == "cross":
                items = [i for i in items if i["category"] == "cross"]
            elif kind == "detect":
                items = [i for i in items if i["category"] == "detect"]
            elif kind in ("in", "out"):
                items = [i for i in items if i["kind"] == kind]
        if track_id is not None:
            items = [i for i in items if i["track_id"] == track_id]
        if date:
            items = [i for i in items if i["timestamp"][:10] == date]

        total = len(items)
        items = items[:limit]
        for i in items:
            i.pop("mtime", None)
        return jsonify({"success": True, "total": total, "count": len(items), "snapshots": items})
    except Exception as e:
        return jsonify({"success": False, "error": str(e), "snapshots": []}), 200


@app.route("/api/snapshot-image/<category>/<path:filename>")
def api_snapshot_image(category, filename):
    if category not in ("cross", "detect"):
        return jsonify({"success": False, "error": "invalid category"}), 404
    base = os.path.abspath(os.path.join(CROSS_SNAPSHOT_DIR, category))
    requested = os.path.abspath(os.path.join(base, filename))
    if not requested.startswith(base + os.sep) or not os.path.isfile(requested):
        return jsonify({"success": False, "error": "not found"}), 404
    return send_file(requested, mimetype="image/jpeg")


def _ensure_db():
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.execute(
        """
        CREATE TABLE IF NOT EXISTS daily_counters (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            counting_date TEXT NOT NULL,
            camera_name TEXT NOT NULL,
            object_label TEXT NOT NULL,
            total_count INTEGER NOT NULL DEFAULT 0,
            total_in INTEGER NOT NULL DEFAULT 0,
            total_out INTEGER NOT NULL DEFAULT 0,
            start_time TEXT,
            end_time TEXT,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE(counting_date, camera_name, object_label)
        )
        """
    )
    conn.commit()
    conn.close()

_ensure_db()


def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def get_counting_date(dt=None, cutoff_str=CUTOFF_TIME):
    if dt is None:
        dt = datetime.now()
    cutoff = datetime.strptime(cutoff_str, "%H:%M").time()
    if dt.time() < cutoff:
        return dt.date().isoformat()
    return (dt.date() + timedelta(days=1)).isoformat()


@app.route("/")
def index():
    return render_template(
        "dashboard.html",
        site_name=SITE_NAME,
        show_detect=SAVE_DETECT_SNAPSHOT,
        control_enabled=CONTROL_ENABLED,
    )


def _read_counting_flag():
    try:
        with open(CONTROL_FILE, "r", encoding="utf-8") as f:
            return bool(json.load(f).get("counting", CONTROL_DEFAULT_COUNTING))
    except FileNotFoundError:
        return CONTROL_DEFAULT_COUNTING
    except Exception:
        return CONTROL_DEFAULT_COUNTING


def _write_counting_flag(counting):
    os.makedirs(os.path.dirname(CONTROL_FILE) or ".", exist_ok=True)
    tmp = f"{CONTROL_FILE}.tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump({"counting": bool(counting)}, f)
    os.replace(tmp, CONTROL_FILE)


@app.route("/api/control", methods=["GET"])
def api_control_get():
    return jsonify({
        "success": True,
        "enabled": CONTROL_ENABLED,
        "counting": _read_counting_flag(),
    })


@app.route("/api/control", methods=["POST"])
def api_control_set():
    if not CONTROL_ENABLED:
        return jsonify({"success": False, "error": "Runtime control is disabled (set CONTROL_ENABLED=true)"}), 403
    data = request.get_json(silent=True) or {}
    if "counting" not in data:
        return jsonify({"success": False, "error": "Missing 'counting' field"}), 400
    try:
        counting = bool(data["counting"])
        _write_counting_flag(counting)
        return jsonify({"success": True, "counting": counting})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route("/snapshots")
def snapshots_page():
    return render_template("snapshots.html", site_name=SITE_NAME, show_detect=SAVE_DETECT_SNAPSHOT)


@app.route("/api/current-counter")
def api_current_counter():
    try:
        with open(CURRENT_COUNTER_PATH, "r") as f:
            data = json.load(f)
        return jsonify(
            {
                "success": True,
                "counting_date": data.get("counting_date"),
                "count": data.get("count", 0),
                "count_in": data.get("count_in", 0),
                "count_out": data.get("count_out", 0),
                "start_time": data.get("start_time"),
                "last_detection_time": data.get("last_detection_time"),
            }
        )
    except FileNotFoundError:
        return jsonify(
            {
                "success": False,
                "error": "No active counter",
                "count": 0,
                "count_in": 0,
                "count_out": 0,
                "counting_date": None,
            }
        ), 200
    except Exception as e:
        return jsonify(
            {
                "success": False,
                "error": str(e),
                "count": 0,
                "count_in": 0,
                "count_out": 0,
                "counting_date": None,
            }
        ), 500


@app.route("/api/summary")
def api_summary():
    try:
        conn = get_db()
        cur = conn.cursor()
        today = get_counting_date()
        cur.execute(
            """
            SELECT COALESCE(total_count, 0) as total_count,
                   COALESCE(total_in, 0) as total_in,
                   COALESCE(total_out, 0) as total_out
            FROM daily_counters
            WHERE counting_date = ?
        """,
            (today,),
        )
        today_row = cur.fetchone()
        yesterday = (datetime.strptime(today, "%Y-%m-%d").date() - timedelta(days=1)).isoformat()
        cur.execute(
            """
            SELECT COALESCE(total_count, 0) as total_count,
                   COALESCE(total_in, 0) as total_in,
                   COALESCE(total_out, 0) as total_out
            FROM daily_counters
            WHERE counting_date = ?
        """,
            (yesterday,),
        )
        yesterday_row = cur.fetchone()
        cur.execute(
            """
            SELECT COALESCE(SUM(total_count), 0) as grand_total,
                   COALESCE(SUM(total_in), 0) as grand_in,
                   COALESCE(SUM(total_out), 0) as grand_out,
                   COUNT(DISTINCT counting_date) as total_days
            FROM daily_counters
        """
        )
        all_time = cur.fetchone()
        cur.execute("SELECT ROUND(AVG(total_count), 1) as avg_per_day FROM daily_counters")
        avg = cur.fetchone()
        cur.execute(
            """
            SELECT counting_date, total_count
            FROM daily_counters
            ORDER BY total_count DESC
            LIMIT 1
        """
        )
        best = cur.fetchone()
        conn.close()
        return jsonify(
            {
                "today": {
                    "date": today,
                    "total_count": today_row["total_count"] if today_row else 0,
                    "total_in": today_row["total_in"] if today_row else 0,
                    "total_out": today_row["total_out"] if today_row else 0,
                },
                "yesterday": {
                    "date": yesterday,
                    "total_count": yesterday_row["total_count"] if yesterday_row else 0,
                    "total_in": yesterday_row["total_in"] if yesterday_row else 0,
                    "total_out": yesterday_row["total_out"] if yesterday_row else 0,
                },
                "all_time": {
                    "grand_total": all_time["grand_total"],
                    "grand_in": all_time["grand_in"],
                    "grand_out": all_time["grand_out"],
                    "total_days": all_time["total_days"],
                },
                "average_per_day": avg["avg_per_day"] or 0,
                "best_day": {
                    "date": best["counting_date"] if best else None,
                    "count": best["total_count"] if best else 0,
                },
            }
        )
    except sqlite3.OperationalError as e:
        return jsonify({"success": False, "error": f"Database unavailable: {e}", "today": {"date": datetime.now().date().isoformat(), "total_count": 0, "total_in": 0, "total_out": 0}, "yesterday": {"date": "", "total_count": 0, "total_in": 0, "total_out": 0}, "all_time": {"grand_total": 0, "grand_in": 0, "grand_out": 0, "total_days": 0}, "average_per_day": 0, "best_day": {"date": None, "count": 0}}), 200
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route("/api/daily-data")
def api_daily_data():
    try:
        days = request.args.get("days", 30, type=int)
        date_from = (datetime.now() - timedelta(days=days)).date().isoformat()
        conn = get_db()
        cur = conn.cursor()
        cur.execute(
            """
            SELECT counting_date, total_count, total_in, total_out
            FROM daily_counters
            WHERE counting_date >= ?
            ORDER BY counting_date ASC
        """,
            (date_from,),
        )
        daily_data = [
            {
                "date": row["counting_date"],
                "total_count": row["total_count"],
                "total_in": row["total_in"],
                "total_out": row["total_out"],
            }
            for row in cur.fetchall()
        ]
        conn.close()
        return jsonify(daily_data)
    except sqlite3.OperationalError:
        return jsonify([]), 200
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route("/api/available-dates")
def api_available_dates():
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute(
            """
            SELECT counting_date, total_count, total_in, total_out, start_time, end_time
            FROM daily_counters
            ORDER BY counting_date DESC
        """
        )
        dates = [
            {
                "date": row["counting_date"],
                "total_count": row["total_count"],
                "total_in": row["total_in"],
                "total_out": row["total_out"],
                "start_time": row["start_time"],
                "end_time": row["end_time"],
            }
            for row in cur.fetchall()
        ]
        conn.close()
        return jsonify(dates)
    except sqlite3.OperationalError:
        return jsonify([]), 200
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


def _excel_response(wb, filename):
    output = BytesIO()
    wb.save(output)
    output.seek(0)
    return Response(
        output.getvalue(),
        mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": f"attachment; filename={filename}"},
    )


def _style_header(ws, cols):
    header_font = Font(bold=True, color="FFFFFF", size=11)
    header_fill = PatternFill(start_color="2F5496", end_color="2F5496", fill_type="solid")
    thin_border = Border(
        left=Side(style="thin"), right=Side(style="thin"),
        top=Side(style="thin"), bottom=Side(style="thin"),
    )
    for col_idx, (col_letter, text) in enumerate(cols, 1):
        cell = ws.cell(row=1, column=col_idx, value=text)
        cell.font = header_font
        cell.fill = header_fill
        cell.alignment = Alignment(horizontal="center")
        cell.border = thin_border
    ws.freeze_panes = "A2"


def _auto_width(ws):
    for col in ws.columns:
        max_len = 0
        col_letter = col[0].column_letter
        for cell in col:
            if cell.value is not None:
                max_len = max(max_len, len(str(cell.value)))
        ws.column_dimensions[col_letter].width = max_len + 4


@app.route("/api/export-daily-csv")
def export_daily_xlsx():
    try:
        days = request.args.get("days", 30, type=int)
        date_from = (datetime.now() - timedelta(days=days)).date().isoformat()
        conn = get_db()
        cur = conn.cursor()
        cur.execute(
            """
            SELECT counting_date, total_count, total_in, total_out, start_time, end_time
            FROM daily_counters
            WHERE counting_date >= ?
            ORDER BY counting_date ASC
        """,
            (date_from,),
        )
        rows = cur.fetchall()
        conn.close()
    except sqlite3.OperationalError as e:
        return jsonify({"success": False, "error": f"Database unavailable: {e}"}), 503
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500

    wb = Workbook()
    ws = wb.active
    ws.title = "Daily Counters"
    _style_header(ws, [("A", "Date"), ("B", "Total"), ("C", "In"), ("D", "Out"), ("E", "Diff (In-Out)"), ("F", "First Count"), ("G", "Last Count")])

    for r_idx, row in enumerate(rows, 2):
        ws.cell(row=r_idx, column=1, value=row["counting_date"])
        ws.cell(row=r_idx, column=2, value=row["total_count"])
        ws.cell(row=r_idx, column=3, value=row["total_in"])
        ws.cell(row=r_idx, column=4, value=row["total_out"])
        ws.cell(row=r_idx, column=5, value=(row["total_in"] or 0) - (row["total_out"] or 0))
        ws.cell(row=r_idx, column=6, value=row["start_time"])
        ws.cell(row=r_idx, column=7, value=row["end_time"])

    _auto_width(ws)
    filename = f"{SITE_NAME}_daily_records_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
    return _excel_response(wb, filename)


if __name__ == "__main__":
    WSGIRequestHandler.protocol_version = "HTTP/1.1"
    print(f"Jetson counter dashboard at http://{DASHBOARD_HOST}:{DASHBOARD_PORT}")
    print(f"DB: {DB_PATH}")
    print(f"State: {CURRENT_COUNTER_PATH}")
    app.run(host=DASHBOARD_HOST, port=DASHBOARD_PORT, debug=FLASK_DEBUG)