import gc
import json
import logging
import os
import threading
import time
import tracemalloc
from collections import defaultdict
from typing import Any, Dict, List, Optional, Set, Tuple

from messagehandlerbuilder import EventManagerMessageHandlerBuilder
from zeromqeventmanager import EventManager

TCP_STATES = {
    "01": "ESTABLISHED",
    "02": "SYN_SENT",
    "03": "SYN_RECV",
    "04": "FIN_WAIT1",
    "05": "FIN_WAIT2",
    "06": "TIME_WAIT",
    "07": "CLOSE",
    "08": "CLOSE_WAIT",
    "09": "LAST_ACK",
    "0A": "LISTEN",
    "0B": "CLOSING",
}


def _parse_ipv4_addr(hex_addr: str) -> Tuple[str, int]:
    addr_hex, port_hex = hex_addr.split(":")
    port = int(port_hex, 16)
    addr_int = int(addr_hex, 16)
    ip = "{}.{}.{}.{}".format(
        addr_int & 0xFF,
        (addr_int >> 8) & 0xFF,
        (addr_int >> 16) & 0xFF,
        (addr_int >> 24) & 0xFF,
    )
    return ip, port


def _parse_ipv6_addr(hex_addr: str) -> Tuple[str, int]:
    addr_hex, port_hex = hex_addr.split(":")
    port = int(port_hex, 16)

    if len(addr_hex) != 32:
        return addr_hex, port

    groups = [addr_hex[i:i + 8] for i in range(0, 32, 8)]
    reversed_groups = []
    for g in groups:
        b = bytes.fromhex(g)
        reversed_groups.append(b[::-1].hex())

    full_hex = "".join(reversed_groups)

    if full_hex == "00000000000000000000ffff" + full_hex[24:]:
        ipv4_hex = full_hex[24:]
        octets = [int(ipv4_hex[i:i + 2], 16) for i in range(0, 8, 2)]
        ip = "::ffff:{}.{}.{}.{}".format(*octets)
    elif full_hex == "0" * 32:
        ip = "::"
    elif full_hex == "0" * 31 + "1":
        ip = "::1"
    else:
        ipv6_parts = [full_hex[i:i + 4] for i in range(0, 32, 4)]
        ip = ":".join(ipv6_parts)

    return ip, port


class DiagnosticMonitor:
    def __init__(
        self,
        message_handler_builder: EventManagerMessageHandlerBuilder,
        interval_seconds: int = 1800,
        enable_tracemalloc: bool = False,
        tracemalloc_top_n: int = 15,
    ) -> None:
        self.message_handler_builder = message_handler_builder
        self._interval = interval_seconds
        self._enable_tracemalloc = enable_tracemalloc
        self._tracemalloc_top_n = tracemalloc_top_n

        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._seq = 0
        self._start_time = 0.0
        self._prev: Optional[Dict[str, Any]] = None
        self._prev_tracemalloc = None
        self._prev_thread_names: Set[str] = set()
        self._cumulative_thread_churn = 0

        self._logger = logging.getLogger("EventManagerDiagnosticMonitor")

        if self._enable_tracemalloc and not tracemalloc.is_tracing():
            tracemalloc.start(25)

    def start(self) -> None:
        if self._running:
            return
        self._running = True
        self._start_time = time.time()
        self._thread = threading.Thread(
            target=self._loop,
            name="DiagnosticMonitor",
            daemon=True,
        )
        self._thread.start()
        self._logger.warning("[DIAG] Monitor started - interval=%ds tracemalloc=%s", self._interval, self._enable_tracemalloc)

    def stop(self) -> None:
        self._running = False
        if self._thread is not None:
            self._thread.join(timeout=10)

    def _loop(self) -> None:
        self._wait(60)
        while self._running:
            try:
                self._snapshot()
            except Exception:
                self._logger.exception("[DIAG] Snapshot failed")
            self._wait(self._interval)

    def _wait(self, seconds: int) -> None:
        end = time.time() + seconds
        while self._running and time.time() < end:
            time.sleep(1)

    def _snapshot(self) -> None:
        self._seq += 1
        uptime_h = round((time.time() - self._start_time) / 3600, 2)

        snap: Dict[str, Any] = {
            "seq": self._seq,
            "uptime_h": uptime_h,
            "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        }

        snap["mem"] = self._collect_memory()
        snap["fd"] = self._collect_fd()
        snap["tcp"] = self._collect_tcp()
        snap["mhb"] = self._collect_mhb()
        snap["threads"] = self._collect_threads()
        snap["locks"] = self._collect_locks()
        snap["zmq_objects"] = self._collect_zmq_objects()
        snap["gc"] = self._collect_gc()

        if self._enable_tracemalloc:
            snap["tracemalloc"] = self._collect_tracemalloc()

        if self._prev is not None:
            snap["delta"] = self._compute_delta(snap, self._prev)

        self._log(snap)

        snap.get("tcp", {}).pop("details", None)
        self._prev = snap

    # -- collectors --

    def _collect_memory(self) -> Dict[str, Any]:
        result: Dict[str, Any] = {}
        try:
            with open("/proc/self/status", "r") as f:
                for line in f:
                    key = line.split(":")[0]
                    if key in ("VmRSS", "VmSize", "VmPeak", "VmHWM", "RssAnon", "RssFile"):
                        result[key] = int(line.split()[1])
        except (FileNotFoundError, PermissionError):
            result["error"] = "proc not available"
        return result

    def _collect_fd(self) -> Dict[str, Any]:
        result: Dict[str, Any] = {"total": 0, "by_type": {}}
        try:
            fd_dir = "/proc/self/fd"
            fds = os.listdir(fd_dir)
            result["total"] = len(fds)

            counts: Dict[str, int] = defaultdict(int)
            files: List[str] = []

            for fd in fds:
                try:
                    target = os.readlink(os.path.join(fd_dir, fd))
                    if target.startswith("socket:"):
                        counts["socket"] += 1
                    elif target.startswith("pipe:"):
                        counts["pipe"] += 1
                    elif "eventfd" in target:
                        counts["eventfd"] += 1
                    elif "eventpoll" in target:
                        counts["eventpoll"] += 1
                    elif target.startswith("anon_inode:"):
                        counts["anon_inode"] += 1
                    else:
                        counts["file"] += 1
                        if len(files) < 20:
                            files.append(target)
                except OSError:
                    counts["unknown"] += 1

            result["by_type"] = dict(counts)
            if files:
                result["files"] = files
        except (FileNotFoundError, PermissionError):
            result["error"] = "proc not available"
        return result

    def _collect_mhb(self) -> Dict[str, Any]:
        message_handler_builder = self.message_handler_builder
        if message_handler_builder is None:
            return {"not_set": True}

        event_manager: EventManager = message_handler_builder.event_manager

        # Sync subscriptions: subject -> list of subscribed services
        sync_data = {}
        for subject, subscription in event_manager.sync_subscriptions.items():
            sync_data[subject] = {
                "services": len(subscription.services),
                "service_names": list(subscription.get_subscribed_services()),
            }

        # Async subscriptions by service name (one StreamerDevice per service)
        service_name_data = {}
        for service_name, subscription in event_manager.service_name_async_subscriptions.items():
            service_name_data[service_name] = {
                "subjects": len(subscription.subjects),
                "sub_ids": len(subscription.sub_id_last_seen),
                "pull_address": subscription.publish_pull_address,
            }

        # Async subscriptions by sub_id (maps sub_id -> service's AsyncSubscription)
        sub_id_data = {}
        for sub_id, subscription in event_manager.sub_id_async_subscription.items():
            sub_id_data[sub_id] = {
                "service_name": subscription.service_name,
                "subjects": len(subscription.sub_id_subjects.get(sub_id, {})),
            }

        return {
            "sync": {
                "total_subjects": len(sync_data),
                "details": sync_data,
            },
            "async": {
                "service_name_async_subscriptions": {
                    "total_services": len(service_name_data),
                    "details": service_name_data,
                },
                "sub_id_async_subscription": {
                    "total_sub_ids": len(sub_id_data),
                    "details": sub_id_data,
                },
            },
        }

    def _collect_tcp(self) -> Dict[str, Any]:
        by_state: Dict[str, int] = defaultdict(int)
        details: List[Dict[str, Any]] = []
        total_tx = 0
        total_rx = 0

        for tcp_path in ("/proc/self/net/tcp", "/proc/self/net/tcp6"):
            is_v6 = "tcp6" in tcp_path
            try:
                with open(tcp_path, "r") as f:
                    for line in f.readlines()[1:]:
                        parts = line.strip().split()
                        if len(parts) < 10:
                            continue

                        state = TCP_STATES.get(parts[3], "UNKNOWN({})".format(parts[3]))
                        by_state[state] += 1

                        tx_hex, rx_hex = parts[4].split(":")
                        tx = int(tx_hex, 16)
                        rx = int(rx_hex, 16)
                        total_tx += tx
                        total_rx += rx

                        parse = _parse_ipv6_addr if is_v6 else _parse_ipv4_addr
                        local_ip, local_port = parse(parts[1])
                        remote_ip, remote_port = parse(parts[2])

                        entry: Dict[str, Any] = {
                            "local": "{}:{}".format(local_ip, local_port),
                            "remote": "{}:{}".format(remote_ip, remote_port),
                            "state": state,
                            "inode": parts[9],
                        }
                        if tx > 0:
                            entry["tx"] = tx
                        if rx > 0:
                            entry["rx"] = rx

                        details.append(entry)
            except (FileNotFoundError, PermissionError):
                continue

        return {
            "by_state": dict(by_state),
            "details": details,
            "total_tx": total_tx,
            "total_rx": total_rx,
        }

    def _collect_threads(self) -> Dict[str, Any]:
        threads = threading.enumerate()
        groups: Dict[str, int] = defaultdict(int)
        names: List[str] = []

        for t in threads:
            name = t.name
            names.append(name)
            prefix = name.split("_Worker_")[0] if "_Worker_" in name else name
            groups[prefix] += 1

        current_names = set(names)
        new_threads = current_names - self._prev_thread_names
        self._cumulative_thread_churn += len(new_threads)
        self._prev_thread_names = current_names

        return {
            "total": len(threads),
            "by_group": dict(groups),
            "new_since_last": list(new_threads) if self._seq > 1 else [],
            "cumulative_churn": self._cumulative_thread_churn,
        }

    def _collect_locks(self) -> Dict[str, Any]:
        result: Dict[str, Any] = {"total_contexts": 0, "total_locks": 0, "top_contexts": []}
        try:
            from zeromqmessagehandler import HandleEventWorkerThread
            lock = HandleEventWorkerThread.non_reentrant_lock
            locks_dict = HandleEventWorkerThread.non_reentrant_context_locks

            with lock:
                context_sizes = {str(ctx): len(subjects) for ctx, subjects in locks_dict.items()}

            result["total_contexts"] = len(context_sizes)
            result["total_locks"] = sum(context_sizes.values())

            sorted_ctx = sorted(context_sizes.items(), key=lambda x: x[1], reverse=True)
            result["top_contexts"] = sorted_ctx[:10]
        except Exception as e:
            result["error"] = str(e)
        return result

    def _collect_zmq_objects(self) -> Dict[str, Any]:
        zmq_socket_count = 0
        socket_resource_count = 0
        lock_count = 0

        try:
            all_objects = gc.get_objects()
        except MemoryError:
            return {"error": "MemoryError - heap too large for gc.get_objects()"}

        try:
            for obj in all_objects:
                type_name = type(obj).__name__
                if type_name == "Socket":
                    zmq_socket_count += 1
                elif type_name == "SocketResource":
                    socket_resource_count += 1
                elif type_name in ("lock", "_RLock"):
                    lock_count += 1
        finally:
            del all_objects

        return {
            "zmq_socket": zmq_socket_count,
            "socket_resource": socket_resource_count,
            "lock": lock_count,
        }

    def _collect_gc(self) -> Dict[str, Any]:
        stats = gc.get_stats()
        return {
            "garbage": len(gc.garbage),
            "generations": [
                {"coll": s["collections"], "ok": s["collected"], "fail": s["uncollectable"]}
                for s in stats
            ],
        }

    def _collect_tracemalloc(self) -> Dict[str, Any]:
        if not tracemalloc.is_tracing():
            return {"enabled": False}

        snapshot = tracemalloc.take_snapshot()
        top = snapshot.statistics("lineno")

        result: Dict[str, Any] = {
            "enabled": True,
            "total_kb": round(sum(s.size for s in top) / 1024, 1),
            "top": [
                {"loc": str(s.traceback), "kb": round(s.size / 1024, 1), "n": s.count}
                for s in top[:self._tracemalloc_top_n]
            ],
        }

        if self._prev_tracemalloc is not None:
            diff = snapshot.compare_to(self._prev_tracemalloc, "lineno")
            result["growth"] = [
                {"loc": str(s.traceback), "diff_kb": round(s.size_diff / 1024, 1), "n": s.count_diff}
                for s in diff[:10]
                if s.size_diff > 0
            ]

        self._prev_tracemalloc = snapshot
        return result

    # -- delta --

    def _compute_delta(self, curr: Dict[str, Any], prev: Dict[str, Any]) -> Dict[str, Any]:
        d: Dict[str, Any] = {}

        d["fd"] = curr.get("fd", {}).get("total", 0) - prev.get("fd", {}).get("total", 0)
        d["rss_kb"] = curr.get("mem", {}).get("VmRSS", 0) - prev.get("mem", {}).get("VmRSS", 0)
        d["threads"] = curr.get("threads", {}).get("total", 0) - prev.get("threads", {}).get("total", 0)
        d["locks"] = curr.get("locks", {}).get("total_locks", 0) - prev.get("locks", {}).get("total_locks", 0)
        d["zmq_socket"] = curr.get("zmq_objects", {}).get("zmq_socket", 0) - prev.get("zmq_objects", {}).get("zmq_socket", 0)

        curr_states = curr.get("tcp", {}).get("by_state", {})
        prev_states = prev.get("tcp", {}).get("by_state", {})
        state_delta = {}
        for s in set(list(curr_states.keys()) + list(prev_states.keys())):
            diff = curr_states.get(s, 0) - prev_states.get(s, 0)
            if diff != 0:
                state_delta[s] = diff
        if state_delta:
            d["tcp_states"] = state_delta

        fd_curr = curr.get("fd", {}).get("by_type", {})
        fd_prev = prev.get("fd", {}).get("by_type", {})
        fd_delta = {}
        for t in set(list(fd_curr.keys()) + list(fd_prev.keys())):
            diff = fd_curr.get(t, 0) - fd_prev.get(t, 0)
            if diff != 0:
                fd_delta[t] = diff
        if fd_delta:
            d["fd_by_type"] = fd_delta

        return d

    # -- logging --

    def _log(self, snap: Dict[str, Any]) -> None:
        seq = snap["seq"]
        up = snap["uptime_h"]
        mem = snap.get("mem", {})
        fd = snap.get("fd", {})
        tcp = snap.get("tcp", {})
        thr = snap.get("threads", {})
        locks = snap.get("locks", {})
        zmq_obj = snap.get("zmq_objects", {})

        rss_mb = mem.get("VmRSS", 0) / 1024
        vms_mb = mem.get("VmSize", 0) / 1024
        peak_mb = mem.get("VmHWM", 0) / 1024

        self._logger.warning(
            "[DIAG] seq=%d uptime=%.1fh rss=%.1fMB(peak=%.1fMB) vms=%.1fMB "
            "fd=%d(sock=%d,pipe=%d,evfd=%d) threads=%d(churn=%d) "
            "locks=%d/%dctx zmq_sock=%d sock_res=%d",
            seq, up, rss_mb, peak_mb, vms_mb,
            fd.get("total", 0),
            fd.get("by_type", {}).get("socket", 0),
            fd.get("by_type", {}).get("pipe", 0),
            fd.get("by_type", {}).get("eventfd", 0),
            thr.get("total", 0),
            thr.get("cumulative_churn", 0),
            locks.get("total_locks", 0),
            locks.get("total_contexts", 0),
            zmq_obj.get("zmq_socket", 0),
            zmq_obj.get("socket_resource", 0),
        )

        self._log_tcp_sockets(seq, tcp)
        self._log_thread_groups(seq, thr)
        self._log_lock_top(seq, locks)
        self._log_delta(seq, snap.get("delta"))
        self._log_gc(seq, snap.get("gc", {}))

        if "tracemalloc" in snap:
            self._log_tracemalloc(seq, snap["tracemalloc"])

        json_snap = dict(snap)
        tcp_copy = dict(json_snap.get("tcp", {}))
        tcp_copy.pop("details", None)
        json_snap["tcp"] = tcp_copy
        self._logger.info("[DIAG:JSON] seq=%d %s", seq, json.dumps(json_snap, default=str, separators=(",", ":")))

    def _log_tcp_sockets(self, seq: int, tcp: Dict[str, Any]) -> None:
        by_state = tcp.get("by_state", {})
        if not by_state:
            return

        state_str = " ".join("{}={}".format(k, v) for k, v in sorted(by_state.items()))
        self._logger.warning(
            "[DIAG:TCP] seq=%d %s tx_queue=%dB rx_queue=%dB",
            seq, state_str, tcp.get("total_tx", 0), tcp.get("total_rx", 0),
        )

        details = tcp.get("details", [])

        listen = [d["local"] for d in details if d["state"] == "LISTEN"]
        if listen:
            self._logger.warning("[DIAG:LISTEN] seq=%d %s", seq, " ".join(sorted(set(listen))))

        conn_groups: Dict[str, int] = defaultdict(int)
        for d in details:
            if d["state"] == "LISTEN":
                continue
            key = "->{}({})".format(d["remote"], d["state"][:3])
            conn_groups[key] += 1

        if conn_groups:
            sorted_conns = sorted(conn_groups.items(), key=lambda x: -x[1])
            parts = ["{}x{}".format(k, v) for k, v in sorted_conns[:30]]
            self._logger.warning("[DIAG:CONN] seq=%d %s", seq, " ".join(parts))

    def _log_thread_groups(self, seq: int, thr: Dict[str, Any]) -> None:
        groups = thr.get("by_group", {})
        if groups:
            group_str = " ".join("{}={}".format(k, v) for k, v in sorted(groups.items()))
            self._logger.warning("[DIAG:THREADS] seq=%d %s", seq, group_str)

        new = thr.get("new_since_last", [])
        if new:
            self._logger.warning("[DIAG:NEW_THREADS] seq=%d count=%d names=%s", seq, len(new), ",".join(sorted(new)[:20]))

    def _log_lock_top(self, seq: int, locks: Dict[str, Any]) -> None:
        top = locks.get("top_contexts", [])
        if not top:
            return
        parts = ["{}={}".format(ctx, count) for ctx, count in top[:5]]
        self._logger.warning("[DIAG:LOCKS_TOP] seq=%d %s", seq, " ".join(parts))

    def _log_delta(self, seq: int, delta: Optional[Dict[str, Any]]) -> None:
        if delta is None:
            return

        simple = []
        for key in ("fd", "rss_kb", "threads", "locks", "zmq_socket"):
            val = delta.get(key, 0)
            if val != 0:
                simple.append("{}={:+d}".format(key, val))

        tcp_d = delta.get("tcp_states", {})
        for state, val in sorted(tcp_d.items()):
            simple.append("{}={:+d}".format(state, val))

        fd_d = delta.get("fd_by_type", {})
        for fd_type, val in sorted(fd_d.items()):
            simple.append("fd_{}={:+d}".format(fd_type, val))

        if simple:
            self._logger.warning("[DIAG:DELTA] seq=%d %s", seq, " ".join(simple))

    def _log_gc(self, seq: int, gc_info: Dict[str, Any]) -> None:
        garbage = gc_info.get("garbage", 0)
        gens = gc_info.get("generations", [])
        if not gens:
            return

        gen_parts = []
        for i, g in enumerate(gens):
            gen_parts.append("gen{}={}/{}/{}".format(i, g["coll"], g["ok"], g["fail"]))

        self._logger.info("[DIAG:GC] seq=%d garbage=%d %s", seq, garbage, " ".join(gen_parts))

    def _log_tracemalloc(self, seq: int, tm: Dict[str, Any]) -> None:
        if not tm.get("enabled"):
            return

        self._logger.warning("[DIAG:TRACEMALLOC] seq=%d total=%.1fKB", seq, tm.get("total_kb", 0))

        for alloc in tm.get("top", [])[:5]:
            self._logger.info(
                "[DIAG:TRACEMALLOC_TOP] seq=%d %.1fKB n=%d %s",
                seq, alloc["kb"], alloc["n"], alloc["loc"],
            )

        for g in tm.get("growth", [])[:5]:
            self._logger.warning(
                "[DIAG:TRACEMALLOC_GROWTH] seq=%d +%.1fKB n=%+d %s",
                seq, g["diff_kb"], g["n"], g["loc"],
            )
