#!/usr/bin/env python3
import json
import os
import shutil
import time
from pathlib import Path
from datetime import datetime

import paho.mqtt.client as mqtt


MQTT_HOST = os.getenv("TESVOR_MQTT_HOST", "192.168.88.56")
MQTT_PORT = int(os.getenv("TESVOR_MQTT_PORT", "1883"))
MQTT_USER = os.getenv("TESVOR_MQTT_USER", "mqtt")
MQTT_PASSWORD = os.getenv("TESVOR_MQTT_PASSWORD", "5905*1-8")

BASE_DIR = Path(os.getenv("TESVOR_BASE_DIR", "/volume2/docker/tesvor-map"))
PUBLIC_DIR = Path(os.getenv("TESVOR_PUBLIC_DIR", "/volume2/docker/mh-map-gateway/html"))
HA_PUBLIC_DIR = Path(os.getenv("TESVOR_HA_PUBLIC_DIR", "/volume2/docker/homeassistant/www/tesvor-map"))

CURRENT_FILE = BASE_DIR / "current.json"
ARCHIVE_DIR = BASE_DIR / "archive"
ARCHIVE_INDEX_FILE = BASE_DIR / "archive_index.json"

PUBLIC_CURRENT_FILE = PUBLIC_DIR / "current.json"
PUBLIC_ARCHIVE_DIR = PUBLIC_DIR / "archive"
PUBLIC_ARCHIVE_INDEX_FILE = PUBLIC_DIR / "archive_index.json"

HA_PUBLIC_CURRENT_FILE = HA_PUBLIC_DIR / "current.json"
HA_PUBLIC_ARCHIVE_DIR = HA_PUBLIC_DIR / "archive"
HA_PUBLIC_ARCHIVE_INDEX_FILE = HA_PUBLIC_DIR / "archive_index.json"

# Old/custom topics from earlier firmware variants
CUSTOM_POINT_TOPIC = "tesvor/x500/map/points"
CUSTOM_STATE_TOPIC = "tesvor/x500/state"

# Current ESPHome MQTT topics
ESPHOME_MAP_TOPIC = "tesvor/x500/sensor/map_path_json/state"
ESPHOME_STATE_TOPIC = "tesvor/x500/sensor/state/state"
ESPHOME_VACUUM_STATE_TOPIC = "tesvor/x500/sensor/vacuum_state/state"
ESPHOME_AVAILABILITY_TOPIC = "tesvor/x500/status"
RESET_TOPIC = "tesvor/x500/map/reset"

POINT_TOPICS = {
    CUSTOM_POINT_TOPIC,
    ESPHOME_MAP_TOPIC,
}

STATE_TOPICS = {
    CUSTOM_STATE_TOPIC,
    ESPHOME_STATE_TOPIC,
    ESPHOME_VACUUM_STATE_TOPIC,
}

CLEANING_STATES = {
    "cleaning",
    "spot_cleaning",
    "edge_cleaning",
    "zmode_cleaning",
}

FINISHED_STATES = {
    "charging",
    "docked",
}

MIN_ARCHIVE_POINTS = 5

points_by_seq = {}
session_started = None
last_state = None
last_availability = None
archived_this_session = False


def now_iso():
    return datetime.now().isoformat(timespec="seconds")


def ensure_dirs():
    BASE_DIR.mkdir(parents=True, exist_ok=True)
    ARCHIVE_DIR.mkdir(parents=True, exist_ok=True)
    PUBLIC_DIR.mkdir(parents=True, exist_ok=True)
    PUBLIC_ARCHIVE_DIR.mkdir(parents=True, exist_ok=True)
    HA_PUBLIC_DIR.mkdir(parents=True, exist_ok=True)
    HA_PUBLIC_ARCHIVE_DIR.mkdir(parents=True, exist_ok=True)


def atomic_write_json(path: Path, payload):
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(
        json.dumps(payload, ensure_ascii=False, separators=(",", ":")),
        encoding="utf-8",
    )
    tmp.replace(path)
    try:
        os.chmod(path, 0o644)
    except Exception:
        pass


def chmod_public_tree():
    try:
        for root, dirs, files in os.walk(PUBLIC_DIR):
            for d in dirs:
                try:
                    os.chmod(Path(root) / d, 0o755)
                except Exception:
                    pass
            for f in files:
                try:
                    os.chmod(Path(root) / f, 0o644)
                except Exception:
                    pass
    except Exception as e:
        print(f"[WARN] chmod public tree failed: {e}", flush=True)


def build_current_payload():
    ordered = [points_by_seq[k] for k in sorted(points_by_seq.keys())]
    return {
        "device": "tesvor_x500",
        "updated": now_iso(),
        "started": session_started,
        "point_count": len(ordered),
        "points": ordered,
    }


def save_current():
    payload = build_current_payload()
    atomic_write_json(CURRENT_FILE, payload)
    atomic_write_json(PUBLIC_CURRENT_FILE, payload)
    atomic_write_json(HA_PUBLIC_CURRENT_FILE, payload)
    print(f"[SAVE] current point_count={payload['point_count']}", flush=True)


def build_archive_index_payload():
    files = []

    for f in sorted(ARCHIVE_DIR.glob("*.json"), reverse=True):
        try:
            stat = f.stat()
            files.append({
                "name": f.name,
                "path": f"archive/{f.name}",
                "size": stat.st_size,
                "modified": datetime.fromtimestamp(stat.st_mtime).isoformat(timespec="seconds"),
            })
        except Exception as e:
            print(f"[WARN] archive index skip {f}: {e}", flush=True)

    return {
        "updated": now_iso(),
        "files": files,
    }


def write_archive_index():
    payload = build_archive_index_payload()
    atomic_write_json(ARCHIVE_INDEX_FILE, payload)
    atomic_write_json(PUBLIC_ARCHIVE_INDEX_FILE, payload)
    atomic_write_json(HA_PUBLIC_ARCHIVE_INDEX_FILE, payload)
    chmod_public_tree()
    print(f"[INDEX] {len(payload['files'])} archive files", flush=True)


def sync_existing_archives_to_public():
    copied = 0

    for src in ARCHIVE_DIR.glob("*.json"):
        dst = PUBLIC_ARCHIVE_DIR / src.name
        ha_dst = HA_PUBLIC_ARCHIVE_DIR / src.name
        try:
            if not dst.exists() or src.stat().st_size != dst.stat().st_size or src.stat().st_mtime > dst.stat().st_mtime:
                shutil.copy2(src, dst)
                os.chmod(dst, 0o644)
                shutil.copy2(src, ha_dst)
                os.chmod(ha_dst, 0o644)
                copied += 1
        except Exception as e:
            print(f"[WARN] sync archive failed {src}: {e}", flush=True)

    if CURRENT_FILE.exists():
        try:
            shutil.copy2(CURRENT_FILE, PUBLIC_CURRENT_FILE)
            os.chmod(PUBLIC_CURRENT_FILE, 0o644)
            shutil.copy2(CURRENT_FILE, HA_PUBLIC_CURRENT_FILE)
            os.chmod(HA_PUBLIC_CURRENT_FILE, 0o644)
        except Exception as e:
            print(f"[WARN] sync current failed: {e}", flush=True)

    if copied:
        print(f"[SYNC] copied {copied} archive file(s) to public", flush=True)

    write_archive_index()


def archive_current(reason):
    global archived_this_session

    if archived_this_session:
        return

    if not CURRENT_FILE.exists():
        print("[ARCHIVE] skipped: no current.json", flush=True)
        return

    if len(points_by_seq) < MIN_ARCHIVE_POINTS:
        print(f"[ARCHIVE] skipped: only {len(points_by_seq)} point(s)", flush=True)
        return

    safe_reason = "".join(c if c.isalnum() or c in ("-", "_") else "_" for c in reason)
    ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

    archive_file = ARCHIVE_DIR / f"tesvor_map_{ts}_{safe_reason}.json"
    public_archive_file = PUBLIC_ARCHIVE_DIR / archive_file.name
    ha_public_archive_file = HA_PUBLIC_ARCHIVE_DIR / archive_file.name

    shutil.copy2(CURRENT_FILE, archive_file)
    os.chmod(archive_file, 0o644)

    shutil.copy2(CURRENT_FILE, public_archive_file)
    os.chmod(public_archive_file, 0o644)

    shutil.copy2(CURRENT_FILE, ha_public_archive_file)
    os.chmod(ha_public_archive_file, 0o644)

    archived_this_session = True
    write_archive_index()

    print(f"[ARCHIVE] {archive_file}", flush=True)
    print(f"[PUBLIC]  {public_archive_file}", flush=True)


def reset_session():
    global points_by_seq, session_started, archived_this_session

    points_by_seq = {}
    session_started = now_iso()
    archived_this_session = False

    save_current()
    print("[SESSION] reset", flush=True)


def normalize_state(text):
    return str(text).strip().lower().replace(" ", "_")


def handle_state(state):
    global last_state, session_started

    state = normalize_state(state)
    if not state:
        return

    if state != last_state:
        print(f"[STATE] {state}", flush=True)
        last_state = state

    # Do NOT auto-reset on cleaning.
    # The current map must remain visible until the explicit Map Reset button is pressed.
    if state in CLEANING_STATES:
        if session_started is None:
            session_started = now_iso()
            print("[SESSION] started", flush=True)

    # Archive when robot reaches dock/charging, but keep current.json unchanged.
    if state in FINISHED_STATES:
        archive_current(state)

    save_current()


def handle_map_payload(text, retained=False):
    global session_started

    # Important:
    # retained map_path_json may be an old map. Do not seed a new session from retained map data.
    if retained:
        print("[MAP] retained payload ignored", flush=True)
        return

    if session_started is None:
        reset_session()

    try:
        data = json.loads(text)
    except Exception as e:
        print(f"[WARN] invalid map json: {e}: {text[:200]}", flush=True)
        return

    p_list = data.get("p", [])

    # ESPHome Map Reset publishes {"n":0,"p":[]}.
    # Treat this as a hard reset for the writer state and current.json.
    if data.get("n") == 0 and p_list == []:
        print("[RESET] map_path_json n=0 p=[]", flush=True)
        reset_session()
        return

    if not isinstance(p_list, list):
        return

    added = 0

    # Format A: custom topic
    # {"p":[[seq,x,y,type],...]}
    #
    # Format B: ESPHome text_sensor map_path_json
    # {"n":9255,"p":[[x,y,type],...]}
    total_n = data.get("n")
    use_synth_seq = False

    if p_list and isinstance(p_list[0], list) and len(p_list[0]) == 3:
        use_synth_seq = True
        try:
            total_n = int(total_n)
            start_seq = max(1, total_n - len(p_list) + 1)
        except Exception:
            start_seq = max(points_by_seq.keys(), default=0) + 1

    for idx, p in enumerate(p_list):
        if not isinstance(p, list):
            continue

        try:
            if len(p) == 4:
                seq, x, y, kind = p
                seq = int(seq)
            elif len(p) == 3 and use_synth_seq:
                x, y, kind = p
                seq = int(start_seq + idx)
            else:
                continue

            x = int(x)
            y = int(y)
            kind = int(kind)
        except Exception:
            continue

        if seq not in points_by_seq:
            points_by_seq[seq] = {
                "seq": seq,
                "x": x,
                "y": y,
                "type": kind,
            }
            added += 1

    if added:
        print(f"[POINTS] +{added}, total={len(points_by_seq)}", flush=True)
        save_current()


def on_connect(client, userdata, flags, rc):
    print(f"[MQTT] connected rc={rc}", flush=True)

    for topic in sorted(POINT_TOPICS | STATE_TOPICS | {ESPHOME_AVAILABILITY_TOPIC, RESET_TOPIC}):
        client.subscribe(topic)
        print(f"[MQTT] subscribed {topic}", flush=True)


def on_disconnect(client, userdata, rc):
    print(f"[MQTT] disconnected rc={rc}", flush=True)


def on_message(client, userdata, msg):
    global last_availability

    topic = msg.topic
    text = msg.payload.decode("utf-8", errors="replace").strip()
    retained = bool(getattr(msg, "retain", False))

    if topic == RESET_TOPIC:
        print("[RESET] mqtt reset requested", flush=True)
        reset_session()
        return

    if topic == ESPHOME_AVAILABILITY_TOPIC:
        availability = normalize_state(text)
        if availability != last_availability:
            print(f"[AVAIL] {availability}", flush=True)
            last_availability = availability
        return

    if topic in STATE_TOPICS:
        if retained:
            print(f"[STATE] retained ignored: {text}", flush=True)
            return
        handle_state(text)
        return

    if topic in POINT_TOPICS:
        handle_map_payload(text, retained=retained)
        return


def main():
    ensure_dirs()
    sync_existing_archives_to_public()

    print("[CONFIG]", flush=True)
    print(f"  MQTT:        {MQTT_HOST}:{MQTT_PORT}", flush=True)
    print(f"  BASE_DIR:    {BASE_DIR}", flush=True)
    print(f"  PUBLIC_DIR:  {PUBLIC_DIR}", flush=True)
    print(f"  HA_PUBLIC_DIR: {HA_PUBLIC_DIR}", flush=True)
    print(f"  POINT_TOPICS: {', '.join(sorted(POINT_TOPICS))}", flush=True)
    print(f"  STATE_TOPICS: {', '.join(sorted(STATE_TOPICS))}", flush=True)

    client = mqtt.Client()

    if MQTT_USER:
        client.username_pw_set(MQTT_USER, MQTT_PASSWORD)

    client.on_connect = on_connect
    client.on_disconnect = on_disconnect
    client.on_message = on_message

    while True:
        try:
            client.connect(MQTT_HOST, MQTT_PORT, 60)
            client.loop_forever()
        except KeyboardInterrupt:
            print("[EXIT] interrupted", flush=True)
            break
        except Exception as e:
            print(f"[ERROR] mqtt loop failed: {e}", flush=True)
            time.sleep(10)


if __name__ == "__main__":
    main()
EOF