#!/usr/bin/env python3
import ipaddress
import json
import os
import socket
import struct
import threading
import time
import concurrent.futures
from collections import OrderedDict

try:
    import tomllib
except ModuleNotFoundError:
    import tomli as tomllib

import requests
from flask import Flask, jsonify, render_template, request

CTRL_PORT = 50742
DISCOVERY_PORT = 50744
BOARD_PORT = 50743
SCAN_TIMEOUT = 1.5
POLL_INTERVAL = 3

app = Flask(__name__)

boards = OrderedDict()          # display_id -> board record
boards_lock = threading.Lock()
ip_index = {}                   # ip -> display_id
ip_identity = {}                # ip -> {"id": identity, "name": display_name}
offline_threshold = POLL_INTERVAL * 3

conflicts = OrderedDict()       # display_id -> conflict record
conflicts_lock = threading.Lock()


def get_local_ips():
    ips = set()
    try:
        _, _, ip_list = socket.gethostbyname_ex(socket.gethostname())
        ips.update(ip_list)
    except:
        pass
    try:
        for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
            addr = info[4][0]
            if addr != "127.0.0.1":
                ips.add(addr)
    except:
        pass
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("10.255.255.255", 1))
        ips.add(s.getsockname()[0])
    except:
        pass
    finally:
        s.close()
    return sorted(ips)


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG_PATH = os.path.join(BASE_DIR, "config.toml")


def load_scan_subnets():
    try:
        with open(CONFIG_PATH, "rb") as f:
            data = tomllib.load(f)
        val = data.get("scan", {}).get("subnets")
        if isinstance(val, str):
            val = [val]
        return val
    except Exception:
        return None


def resolve_subnets(patterns):
    if patterns is None:
        return None
    networks = []
    for pat in patterns:
        pat = str(pat).strip()
        if not pat or pat.lower() == "all":
            return None
        try:
            if "*" in pat:
                prefix = 0
                cidr_parts = []
                for part in pat.split("."):
                    if part == "*":
                        break
                    cidr_parts.append(part)
                    prefix += 8
                if prefix == 0:
                    continue
                padded = cidr_parts + ["0"] * (4 - len(cidr_parts))
                networks.append(ipaddress.ip_network(
                    ".".join(padded) + "/" + str(prefix), strict=False))
            elif "/" in pat:
                networks.append(ipaddress.ip_network(pat, strict=False))
            else:
                networks.append(ipaddress.ip_network(pat + "/32", strict=False))
        except ValueError:
            print(f"  Ignoring invalid subnet pattern: {pat}")
    return networks


_scan_networks_cache = None


def get_scan_networks():
    global _scan_networks_cache
    if _scan_networks_cache is None:
        _scan_networks_cache = resolve_subnets(load_scan_subnets())
    return _scan_networks_cache


def ip_allowed(ip):
    networks = get_scan_networks()
    if networks is None:
        return True
    try:
        addr = ipaddress.ip_address(ip)
    except ValueError:
        return False
    return any(addr in net for net in networks)


def build_scan_targets(networks):
    targets = []
    if networks is None:
        subnets = set()
        for ip in get_local_ips():
            subnets.add(".".join(ip.split(".")[:3]))
        for subnet in sorted(subnets):
            for i in range(1, 255):
                targets.append(f"{subnet}.{i}")
    else:
        for net in networks:
            for ip in net.hosts():
                targets.append(str(ip))
    return targets


def check_port(ip, port, timeout=SCAN_TIMEOUT):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        s.connect((ip, port))
        s.close()
        return ip
    except:
        return None


def fetch_status(ip, port):
    try:
        r = requests.get(f"http://{ip}:{port}/api/status", timeout=3)
        if r.status_code == 200:
            return r.json()
    except:
        pass
    return None


def board_identity(status):
    """Static identity of a board so multi-homed IPs can be told apart
    from distinct physical boards that share a display_id."""
    if not status:
        return None
    cfg = status.get("config") or {}
    return (
        status.get("display_id"),
        status.get("display_name"),
        status.get("header"),
        json.dumps(cfg.get("left"), sort_keys=True, default=str),
        json.dumps(cfg.get("right"), sort_keys=True, default=str),
    )


def record_conflict(did, ip, status, other_ips):
    """Record that multiple physical boards are using the same display_id."""
    now = time.time()
    incoming = board_identity(status)
    with conflicts_lock:
        groups = {}
        for other in set(list(other_ips) + [ip]):
            info = ip_identity.get(other)
            ident = info["id"] if info else None
            name = info["name"] if info else ""
            if other == ip and incoming is not None:
                ident = incoming
                name = (status or {}).get("display_name", "")
            if ident is None:
                continue
            groups.setdefault(ident, []).append({"ip": other, "display_name": name})
        if len(groups) < 2:
            conflicts.pop(did, None)
            return
        rec = {"display_id": did, "boards": [], "last_seen": now}
        for entries in groups.values():
            for e in entries:
                rec["boards"].append({
                    "ip": e["ip"],
                    "display_name": e["display_name"],
                    "last_seen": now,
                })
        conflicts[did] = rec


def register_board(did, ip, port, display_name="", status=None):
    identity = board_identity(status)
    conflict_ips = None
    with boards_lock:
        now = time.time()
        prev_did = ip_index.get(ip)
        if prev_did is not None and prev_did != did and prev_did in boards:
            old = boards[prev_did]
            if ip in old["ips"]:
                old["ips"].remove(ip)
            if not old["ips"]:
                boards.pop(prev_did, None)
        if did in boards:
            board = boards[did]
            if identity is not None:
                for other in list(board["ips"]):
                    other_info = ip_identity.get(other)
                    if other_info and other_info["id"] != identity:
                        conflict_ips = list(board["ips"])
                        break
            if ip not in board["ips"]:
                board["ips"].append(ip)
                print(f"  Board {did} also reachable at {ip}:{port}")
            if display_name and not board["display_name"]:
                board["display_name"] = display_name
            board["port"] = port
            board["last_seen"] = now
            board["online"] = True
            board["active_ip"] = ip
            if status is not None:
                board["status"] = status
                if status.get("display_name"):
                    board["display_name"] = status["display_name"]
            boards.move_to_end(did)
        else:
            boards[did] = {
                "display_id": did,
                "display_name": display_name or f"Board ({ip})",
                "ips": [ip],
                "port": port,
                "last_seen": now,
                "online": True,
                "status": status,
                "active_ip": ip,
            }
            print(f"  Discovered board: {boards[did]['display_name']} at {ip}:{port}")
        if identity is not None:
            ip_identity[ip] = {"id": identity, "name": display_name or (status or {}).get("display_name", "")}
        ip_index[ip] = did
    if conflict_ips is not None:
        record_conflict(did, ip, status, conflict_ips)


def udp_listener():
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.settimeout(1)
    sock.bind(("0.0.0.0", DISCOVERY_PORT))
    while True:
        try:
            data, addr = sock.recvfrom(2048)
            msg = json.loads(data)
            port = msg.get("port", BOARD_PORT)
            did = msg.get("display_id")
            ips = msg.get("ips") or []
            if not ips:
                ips = [msg.get("ip")] if msg.get("ip") else [addr[0]]
            for ip in ips:
                if not ip_allowed(ip):
                    print(f"  Ignoring {ip}: outside scan subnets")
                    continue
                register_board(did or ip, ip, port, msg.get("display_name", ""))
        except socket.timeout:
            pass
        except:
            pass


def scan_subnets():
    patterns = load_scan_subnets()
    networks = resolve_subnets(patterns)
    targets = build_scan_targets(networks)
    if networks is not None:
        print(f"  Scanning subnets: {', '.join(str(n) for n in networks)}")
    else:
        print("  Scanning local subnets (all)")
    print(f"  Scanning {len(targets)} hosts")
    found = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=100) as pool:
        futures = {}
        for target in targets:
            futures[pool.submit(check_port, target, BOARD_PORT)] = target
        for future in concurrent.futures.as_completed(futures):
            ip = future.result()
            if ip:
                found.append(ip)
    print(f"  Scan found IPs: {', '.join(sorted(found))}")
    for ip in found:
        status = fetch_status(ip, BOARD_PORT)
        if status and status.get("display_id"):
            did = status["display_id"]
            with boards_lock:
                known = did in boards
            if known:
                print(f"  Deduplicated: {ip} -> display_id {did} (already registered)")
            register_board(did, ip, BOARD_PORT,
                           status.get("display_name", ""), status)
        else:
            with boards_lock:
                known = ip in ip_index
            if known:
                print(f"  Deduplicated: {ip} (already registered)")
            register_board(ip, ip, BOARD_PORT, f"Board ({ip})", status)
    return found


def refresh_board_cache(ip, port):
    status = fetch_status(ip, port)
    if status is None:
        return
    did = status.get("display_id") or ip
    register_board(did, ip, port, status.get("display_name", ""), status)


def refresh_conflicts():
    """Rebuild conflict records from current ip -> identity data so
    conflicts disappear once display_ids are made unique."""
    now = time.time()
    with boards_lock:
        by_did = {}
        for ip, did in list(ip_index.items()):
            info = ip_identity.get(ip)
            if info:
                by_did.setdefault(did, []).append((ip, info["id"], info["name"]))
    with conflicts_lock:
        for did in list(conflicts.keys()):
            if did not in by_did:
                conflicts.pop(did, None)
        for did, entries in by_did.items():
            idents = {ident for _, ident, _ in entries}
            if len(idents) >= 2:
                conflicts[did] = {
                    "display_id": did,
                    "boards": [{"ip": ip, "display_name": name, "last_seen": now}
                               for ip, _ident, name in entries],
                    "last_seen": now,
                }
            else:
                conflicts.pop(did, None)


def poll_boards():
    while True:
        with boards_lock:
            snapshot = list(boards.values())
        for board in snapshot:
            status = None
            active_ip = None
            for ip in list(board["ips"]):
                status = fetch_status(ip, board["port"])
                if status is None:
                    continue
                reported_did = status.get("display_id")
                if reported_did is not None and reported_did != board["display_id"]:
                    with boards_lock:
                        rec = boards.get(board["display_id"])
                        if rec and ip in rec["ips"]:
                            rec["ips"].remove(ip)
                        if rec and not rec["ips"]:
                            boards.pop(board["display_id"], None)
                    status = None
                    continue
                active_ip = ip
                break
            with boards_lock:
                if board["display_id"] not in boards:
                    continue
                rec = boards[board["display_id"]]
                if status is not None:
                    rec["status"] = status
                    rec["online"] = True
                    rec["last_seen"] = time.time()
                    rec["active_ip"] = active_ip
                    if status.get("display_name"):
                        rec["display_name"] = status["display_name"]
                    if active_ip in rec["ips"]:
                        rec["ips"].remove(active_ip)
                        rec["ips"].insert(0, active_ip)
                else:
                    rec["online"] = False
        with boards_lock:
            now = time.time()
            dead = [did for did, b in list(boards.items())
                    if not b.get("online") and (now - b["last_seen"]) > 600]
            for did in dead:
                boards.pop(did, None)
            for ip in [ip for ip, d in list(ip_index.items()) if d in dead]:
                ip_index.pop(ip, None)
            for ip in [ip for ip in list(ip_identity.keys()) if ip not in ip_index]:
                ip_identity.pop(ip, None)
        refresh_conflicts()
        time.sleep(POLL_INTERVAL)


threading.Thread(target=udp_listener, daemon=True).start()
threading.Thread(target=poll_boards, daemon=True).start()


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/boards")
def api_boards():
    now = time.time()
    with boards_lock:
        result = []
        for did, board in list(boards.items()):
            stale = (now - board["last_seen"]) > offline_threshold
            status = board.get("status") or {}
            scores = status.get("scores") or {}
            config_data = status.get("config") or {}
            result.append({
                "display_id": did,
                "display_name": board["display_name"],
                "ip": board.get("active_ip") or (board["ips"][0] if board["ips"] else None),
                "ips": list(board["ips"]),
                "port": board["port"],
                "online": board["online"] and not stale,
                "scores": scores,
                "config": config_data,
            })
        return jsonify(result)


@app.route("/api/conflicts")
def api_conflicts():
    with conflicts_lock:
        return jsonify(list(conflicts.values()))


@app.route("/api/scan", methods=["POST"])
def api_scan():
    def do_scan():
        found = scan_subnets()
        print(f"  Scan complete: {len(found)} board IPs found")
    threading.Thread(target=do_scan, daemon=True).start()
    return jsonify({"status": "scanning"})


@app.route("/api/boards/<ip>/<path:endpoint>", methods=["GET", "POST", "PUT", "OPTIONS"])
def api_proxy(ip, endpoint):
    port = BOARD_PORT
    with boards_lock:
        did = ip_index.get(ip)
        if did is not None and did in boards:
            board = boards[did]
            port = board["port"]
            if board.get("active_ip"):
                ip = board["active_ip"]
    url = f"http://{ip}:{port}/{endpoint}"
    method = request.method
    headers = {"Content-Type": "application/json"} if request.is_json else {}
    try:
        if method == "GET":
            r = requests.get(url, timeout=5, headers=headers)
        elif method == "POST":
            r = requests.post(url, timeout=5, json=request.get_json(silent=True), headers=headers)
        elif method == "PUT":
            r = requests.put(url, timeout=5, json=request.get_json(silent=True), headers=headers)
        else:
            r = requests.options(url, timeout=5, headers=headers)
        resp = app.make_response((r.content, r.status_code, dict(r.headers)))
        resp.headers["Access-Control-Allow-Origin"] = "*"
        # Immediately refresh cached status after any score-modifying action
        if method in ("POST", "PUT"):
            threading.Thread(target=refresh_board_cache, args=(ip, port), daemon=True).start()
        return resp
    except Exception as e:
        return jsonify({"error": str(e)}), 502


@app.after_request
def add_cors(resp):
    resp.headers["Access-Control-Allow-Origin"] = "*"
    resp.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, DELETE, OPTIONS"
    resp.headers["Access-Control-Allow-Headers"] = "Content-Type"
    return resp


if __name__ == "__main__":
    print(f"  CuteBoard Controller")
    print(f"  URL: http://0.0.0.0:{CTRL_PORT}")
    app.run(host="0.0.0.0", port=CTRL_PORT, debug=False, threaded=True)