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


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,
        interval_seconds: int = 1800,
        enable_tracemalloc: bool = False,
        tracemalloc_top_n: int = 15,
    ) -> None:
        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["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_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"],
            )
