import gzip
from collections import Counter
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
from profile_pb2 import Profile


class _StackTreeNode:
    __slots__ = ("children", "weight_a", "weight_b")

    def __init__(self):
        self.children: Dict[str, "_StackTreeNode"] = {}
        self.weight_a: float = 0.0
        self.weight_b: float = 0.0


@dataclass
class ProfileComparisonResult:
    meta: Dict[str, Any]
    metrics: Dict[str, Any]


class ProfileAnalyzer:
    def __init__(self, top_n: int = 20, stack_top_k: int = 100):
        self.top_n = top_n
        self.stack_top_k = stack_top_k
        self.metric_specs = {
            "samples": ["samples", "count"],
            "cpu": ["cpu", "nanoseconds", "ns"],
        }

    def compare(self, path_a: str, path_b: str) -> ProfileComparisonResult:
        profile_a = self._load_profile(path_a)
        profile_b = self._load_profile(path_b)

        sample_types_a = [
            self._sample_type_descriptor(profile_a, idx)
            for idx in range(len(profile_a.sample_type))
        ]
        sample_types_b = [
            self._sample_type_descriptor(profile_b, idx)
            for idx in range(len(profile_b.sample_type))
        ]

        metrics: Dict[str, Dict[str, Any]] = {}
        for metric_name, keywords in self.metric_specs.items():
            idx_a = self._find_sample_index(profile_a, keywords)
            idx_b = self._find_sample_index(profile_b, keywords)
            if idx_a is None and idx_b is None:
                continue
            if idx_a is None:
                idx_a = idx_b if idx_b is not None else 0
            if idx_b is None:
                idx_b = idx_a

            exclusive_a, inclusive_a = self._build_function_aggregations(
                profile_a, idx_a
            )
            exclusive_b, inclusive_b = self._build_function_aggregations(
                profile_b, idx_b
            )

            flat_topn = self._top_n_similarity(exclusive_a, exclusive_b)
            cum_topn = self._top_n_similarity(inclusive_a, inclusive_b)

            stacks_a = self._build_collapsed_stacks(profile_a, idx_a)
            stacks_b = self._build_collapsed_stacks(profile_b, idx_b)

            metrics[metric_name] = {
                "value_index_a": idx_a,
                "value_index_b": idx_b,
                "value_type_a": (
                    sample_types_a[idx_a] if idx_a < len(sample_types_a) else None
                ),
                "value_type_b": (
                    sample_types_b[idx_b] if idx_b < len(sample_types_b) else None
                ),
                "function_topn": {
                    "flat": flat_topn,
                    "cum": cum_topn,
                },
                "function_global": self._global_distribution_similarity(
                    exclusive_a, exclusive_b
                ),
                "stack_similarity": self._stack_similarity(stacks_a, stacks_b),
                "function_summary": {
                    "profile_a": self._compute_function_tables(
                        exclusive_a, inclusive_a
                    ),
                    "profile_b": self._compute_function_tables(
                        exclusive_b, inclusive_b
                    ),
                },
            }

        meta = {
            "time_a": profile_a.time_nanos,
            "time_b": profile_b.time_nanos,
            "duration_a": profile_a.duration_nanos,
            "duration_b": profile_b.duration_nanos,
            "sample_types_a": sample_types_a,
            "sample_types_b": sample_types_b,
        }
        return ProfileComparisonResult(meta=meta, metrics=metrics)

    @staticmethod
    def _load_profile(path: str) -> Profile:
        with gzip.open(path, "rb") as f:
            data = f.read()
        profile = Profile()
        profile.ParseFromString(data)
        return profile

    def _top_n_similarity(
        self, counts_a: Dict[str, float], counts_b: Dict[str, float]
    ) -> Dict[str, Any]:
        n = self.top_n
        top_a = sorted(counts_a.items(), key=lambda kv: kv[1], reverse=True)[:n]
        top_b = sorted(counts_b.items(), key=lambda kv: kv[1], reverse=True)[:n]
        funcs_a = [name for name, _ in top_a]
        funcs_b = [name for name, _ in top_b]
        jaccard = self._jaccard_similarity(funcs_a, funcs_b)
        union_funcs = sorted(set(funcs_a) | set(funcs_b))
        vec_a = np.array([counts_a.get(name, 0.0) for name in union_funcs], dtype=float)
        vec_b = np.array([counts_b.get(name, 0.0) for name in union_funcs], dtype=float)
        cosine = self._cosine_similarity(vec_a, vec_b)
        return {
            "top_n_jaccard": jaccard,
            "top_n_cosine": cosine,
            "top_n_list_a": top_a,
            "top_n_list_b": top_b,
        }

    def _compute_function_tables(
        self,
        exclusive: Dict[str, float],
        inclusive: Dict[str, float],
    ) -> Dict[str, Any]:
        max_entries = self.top_n
        total_flat = sum(exclusive.values())
        flat_sorted = sorted(exclusive.items(), key=lambda kv: kv[1], reverse=True)
        running = 0.0
        flat_entries = []
        for idx, (func, flat) in enumerate(flat_sorted):
            running += flat
            cum_value = inclusive.get(func, 0.0)
            entry = {
                "function": func,
                "flat": flat,
                "flat_percent": (flat / total_flat) if total_flat else 0.0,
                "sum_percent": (running / total_flat) if total_flat else 0.0,
                "cum": cum_value,
                "cum_percent": (cum_value / total_flat) if total_flat else 0.0,
            }
            if max_entries is None or idx < max_entries:
                flat_entries.append(entry)
        cum_sorted = sorted(inclusive.items(), key=lambda kv: kv[1], reverse=True)
        cum_entries = []
        for idx, (func, cum_value) in enumerate(cum_sorted):
            entry = {
                "function": func,
                "cum": cum_value,
                "cum_percent": (cum_value / total_flat) if total_flat else 0.0,
                "flat": exclusive.get(func, 0.0),
                "flat_percent": (
                    (exclusive.get(func, 0.0) / total_flat) if total_flat else 0.0
                ),
            }
            if max_entries is None or idx < max_entries:
                cum_entries.append(entry)
        return {
            "total_flat": total_flat,
            "flat_entry_count": len(flat_sorted),
            "cum_entry_count": len(cum_sorted),
            "flat_table": flat_entries,
            "cum_table": cum_entries,
        }

    def _stack_similarity(
        self,
        stacks_a: Counter,
        stacks_b: Counter,
    ) -> Dict[str, Any]:
        top_k = self.stack_top_k
        top_a = stacks_a.most_common(top_k)
        top_b = stacks_b.most_common(top_k)
        stacks_a_names = [stack for stack, _ in top_a]
        stacks_b_names = [stack for stack, _ in top_b]
        jaccard = self._jaccard_similarity(stacks_a_names, stacks_b_names)
        union = sorted(set(stacks_a_names) | set(stacks_b_names))
        if union:
            vec_a = np.array([stacks_a.get(stack, 0.0) for stack in union], dtype=float)
            vec_b = np.array([stacks_b.get(stack, 0.0) for stack in union], dtype=float)
            cosine = self._cosine_similarity(vec_a, vec_b)
        else:
            cosine = 0.0
        tree_stats = self._stack_tree_stats(stacks_a, stacks_b)
        return {
            "stack_top_k_cosine": cosine,
            "stack_top_k_jaccard": jaccard,
            "stack_tree_emd": tree_stats["emd"],
            "stack_tree_total_a": tree_stats["total_a"],
            "stack_tree_total_b": tree_stats["total_b"],
            "stack_top_a": top_a,
            "stack_top_b": top_b,
        }

    def _stack_tree_stats(
        self,
        stacks_a: Counter,
        stacks_b: Counter,
    ) -> Dict[str, float]:
        total_a = float(sum(stacks_a.values()))
        total_b = float(sum(stacks_b.values()))
        if total_a == 0.0 and total_b == 0.0:
            return {"emd": 0.0, "total_a": 0.0, "total_b": 0.0}

        def normalize(counter: Counter, total: float) -> Dict[str, float]:
            if total <= 0.0:
                return {}
            return {key: value / total for key, value in counter.items()}

        norm_a = normalize(stacks_a, total_a)
        norm_b = normalize(stacks_b, total_b)
        root = _StackTreeNode()
        for path, weight in norm_a.items():
            self._stack_tree_add(root, path, weight, True)
        for path, weight in norm_b.items():
            self._stack_tree_add(root, path, weight, False)

        emd_cost, root_diff = self._stack_tree_emd(root)
        # root_diff should be ~0 for normalized weights; ignore residual
        return {"emd": emd_cost, "total_a": total_a, "total_b": total_b}

    @staticmethod
    def _stack_tree_add(
        root: _StackTreeNode, path: str, weight: float, is_a: bool
    ) -> None:
        if not path:
            parts = ["<unknown>"]
        else:
            parts = path.split(";")
        node = root
        for part in parts:
            if not part:
                part = "<unknown>"
            if part not in node.children:
                node.children[part] = _StackTreeNode()
            node = node.children[part]
        if is_a:
            node.weight_a += weight
        else:
            node.weight_b += weight

    @staticmethod
    def _stack_tree_emd(node: _StackTreeNode) -> Tuple[float, float]:
        # maybe other more explainable distance metrics. Tree Edit Distance? Weighted Path?
        diff = node.weight_a - node.weight_b
        cost = 0.0
        for child in node.children.values():
            child_cost, child_diff = ProfileAnalyzer._stack_tree_emd(child)
            cost += child_cost + abs(child_diff)
            diff += child_diff
        return cost, diff

    def _global_distribution_similarity(
        self,
        counts_a: Dict[str, float],
        counts_b: Dict[str, float],
    ) -> Dict[str, float]:
        all_funcs = sorted(set(counts_a) | set(counts_b))
        vec_a = np.array([counts_a.get(func, 0.0) for func in all_funcs], dtype=float)
        vec_b = np.array([counts_b.get(func, 0.0) for func in all_funcs], dtype=float)
        total_a = vec_a.sum()
        total_b = vec_b.sum()
        if total_a > 0:
            vec_a /= total_a
        if total_b > 0:
            vec_b /= total_b
        cosine = self._cosine_similarity(vec_a, vec_b)
        js = self._js_divergence(vec_a, vec_b)
        kl_ab = self._kl_divergence(vec_a, vec_b)
        kl_ba = self._kl_divergence(vec_b, vec_a)
        return {
            "cosine": cosine,
            "js_divergence": js,
            "kl_a_to_b": kl_ab,
            "kl_b_to_a": kl_ba,
            "total_a": float(total_a),
            "total_b": float(total_b),
        }

    def _build_function_aggregations(
        self,
        profile: Profile,
        value_index: int = 0,
    ) -> Tuple[Dict[str, float], Dict[str, float]]:
        func_meta = {
            fn.id: self._decode_string(profile, fn.name) or "<unknown>"
            for fn in profile.function
        }
        loc_to_funcs: Dict[int, List[str]] = {}
        loc_to_leaf: Dict[int, str] = {}
        for location in profile.location:
            names = [
                func_meta.get(line.function_id, "<unknown>") for line in location.line
            ]
            if not names:
                names = ["<unknown>"]
            loc_to_funcs[location.id] = names
            loc_to_leaf[location.id] = names[0]
        exclusive: Dict[str, float] = {}
        inclusive: Dict[str, float] = {}
        for sample in profile.sample:
            if not sample.location_id:
                continue
            if value_index < len(sample.value):
                value = float(sample.value[value_index])
            elif sample.value:
                value = float(sample.value[0])
            else:
                value = 0.0
            leaf_loc = sample.location_id[0]
            leaf_name = loc_to_leaf.get(leaf_loc, "<unknown>")
            exclusive[leaf_name] = exclusive.get(leaf_name, 0.0) + value
            for loc_id in sample.location_id:
                for func_name in loc_to_funcs.get(loc_id, ["<unknown>"]):
                    inclusive[func_name] = inclusive.get(func_name, 0.0) + value
        return exclusive, inclusive

    def _build_collapsed_stacks(
        self,
        profile: Profile,
        value_index: int = 0,
    ) -> Counter:
        func_meta = {
            fn.id: self._decode_string(profile, fn.name) or "<unknown>"
            for fn in profile.function
        }
        loc_to_funcs: Dict[int, List[str]] = {}
        for location in profile.location:
            names = [
                func_meta.get(line.function_id, "<unknown>") for line in location.line
            ]
            loc_to_funcs[location.id] = names or ["<unknown>"]
        counter: Counter = Counter()
        for sample in profile.sample:
            stack_funcs: List[str] = []
            for loc_id in sample.location_id:
                stack_funcs.extend(loc_to_funcs.get(loc_id, ["<unknown>"]))
            collapsed = ";".join(stack_funcs)
            if value_index < len(sample.value):
                value = float(sample.value[value_index])
            elif sample.value:
                value = float(sample.value[0])
            else:
                value = 0.0
            counter[collapsed] += value
        return counter

    @staticmethod
    def _sample_type_descriptor(
        profile: Profile,
        idx: int,
    ) -> Optional[Tuple[str, str]]:
        if idx < 0 or idx >= len(profile.sample_type):
            return None
        sample_type = profile.sample_type[idx]
        return (
            ProfileAnalyzer._decode_string(profile, sample_type.type),
            ProfileAnalyzer._decode_string(profile, sample_type.unit),
        )

    @staticmethod
    def _find_sample_index(profile: Profile, keywords: List[str]) -> Optional[int]:
        lowered = {kw.lower() for kw in keywords}
        for idx, sample_type in enumerate(profile.sample_type):
            type_name = ProfileAnalyzer._decode_string(
                profile, sample_type.type
            ).lower()
            unit_name = ProfileAnalyzer._decode_string(
                profile, sample_type.unit
            ).lower()
            if type_name in lowered or unit_name in lowered:
                return idx
        return None

    @staticmethod
    def _decode_string(profile: Profile, idx: int) -> str:
        if idx < 0 or idx >= len(profile.string_table):
            return "<invalid>"
        value = profile.string_table[idx]
        if isinstance(value, bytes):
            return value.decode("utf-8", errors="replace")
        return value or ""

    @staticmethod
    def _cosine_similarity(u: np.ndarray, v: np.ndarray) -> float:
        u = np.asarray(u, dtype=float).ravel()
        v = np.asarray(v, dtype=float).ravel()
        denom = np.linalg.norm(u) * np.linalg.norm(v)
        if denom == 0:
            return 0.0
        return float(u @ v / denom)

    @staticmethod
    def _jaccard_similarity(list1: List[str], list2: List[str]) -> float:
        set1 = set(list1)
        set2 = set(list2)
        union = set1 | set2
        if not union:
            return 0.0
        return len(set1 & set2) / len(union)

    @staticmethod
    def _kl_divergence(p: np.ndarray, q: np.ndarray) -> float:
        eps = 1e-12
        p_safe = p + eps
        q_safe = q + eps
        return float(np.sum(p_safe * np.log(p_safe / q_safe)))

    @staticmethod
    def _js_divergence(p: np.ndarray, q: np.ndarray) -> float:
        m = 0.5 * (p + q)
        return 0.5 * (
            ProfileAnalyzer._kl_divergence(p, m) + ProfileAnalyzer._kl_divergence(q, m)
        )


if __name__ == "__main__":
    import argparse
    import json

    parser = argparse.ArgumentParser(description="Compare two pprof profiles")
    parser.add_argument("profile_a", help="path to profile A (gzipped proto)")
    parser.add_argument("profile_b", help="path to profile B (gzipped proto)")
    parser.add_argument("--top-n", type=int, default=20, dest="top_n")
    parser.add_argument("--stack-top-k", type=int, default=100, dest="stack_top_k")
    parser.add_argument(
        "--output", type=str, default=None, help="optional path to write JSON result"
    )
    parser.add_argument(
        "--metric",
        choices=["samples", "cpu"],
        default="cpu",
        help="metric to summarize (default: cpu)",
    )
    parser.add_argument(
        "--detailed",
        action="store_true",
        help="include detailed top lists and stack entries",
    )
    args = parser.parse_args()

    analyzer = ProfileAnalyzer(top_n=args.top_n, stack_top_k=args.stack_top_k)
    comparison = analyzer.compare(args.profile_a, args.profile_b)

    metrics = comparison.metrics
    metric_key = args.metric
    if metric_key not in metrics:
        available = ", ".join(sorted(metrics.keys()))
        parser.error(
            f"metric '{metric_key}' not available in profile (have: {available})"
        )

    metric_data = metrics[metric_key]

    def serialize_topn(topn: Dict[str, object]) -> Dict[str, object]:
        data = {
            "top_n_jaccard": topn["top_n_jaccard"],
            "top_n_cosine": topn["top_n_cosine"],
        }
        if args.detailed:
            data["top_n_list_a"] = [
                [name, value] for name, value in topn["top_n_list_a"]
            ]
            data["top_n_list_b"] = [
                [name, value] for name, value in topn["top_n_list_b"]
            ]
        return data

    def serialize_stack(stack: Dict[str, object]) -> Dict[str, object]:
        data = {
            "stack_top_k_cosine": stack["stack_top_k_cosine"],
            "stack_top_k_jaccard": stack["stack_top_k_jaccard"],
            "stack_tree_emd": stack.get("stack_tree_emd"),
            "stack_tree_total_a": stack.get("stack_tree_total_a"),
            "stack_tree_total_b": stack.get("stack_tree_total_b"),
        }
        if args.detailed:
            data["stack_top_a"] = [
                [name, value] for name, value in stack["stack_top_a"]
            ]
            data["stack_top_b"] = [
                [name, value] for name, value in stack["stack_top_b"]
            ]
        return data

    function_topn = {
        name: serialize_topn(data)
        for name, data in metric_data["function_topn"].items()
    }

    metric_output = {
        "value_index_a": metric_data["value_index_a"],
        "value_index_b": metric_data["value_index_b"],
        "value_type_a": metric_data["value_type_a"],
        "value_type_b": metric_data["value_type_b"],
        "value_unit": (
            metric_data["value_type_a"][1] if metric_data["value_type_a"] else None
        ),
        "function_topn": function_topn,
        "stack_similarity": serialize_stack(metric_data["stack_similarity"]),
    }

    def unit_to_seconds(unit: Optional[str]) -> Optional[float]:
        if not unit:
            return None
        mapping = {
            "nanoseconds": 1e9,
            "ns": 1e9,
            "microseconds": 1e6,
            "us": 1e6,
            "milliseconds": 1e3,
            "ms": 1e3,
            "seconds": 1.0,
            "s": 1.0,
        }
        return mapping.get(unit.lower())

    def scale_lists(metric_dict: Dict[str, object], scale: float) -> None:
        def scale_pairs(pairs: List[List[object]]) -> List[List[object]]:
            return [[name, value / scale] for name, value in pairs]

        for topn in metric_dict["function_topn"].values():
            if "top_n_list_a" in topn:
                topn["top_n_list_a"] = scale_pairs(topn["top_n_list_a"])
            if "top_n_list_b" in topn:
                topn["top_n_list_b"] = scale_pairs(topn["top_n_list_b"])

        stack = metric_dict["stack_similarity"]
        if "stack_top_a" in stack:
            stack["stack_top_a"] = scale_pairs(stack["stack_top_a"])
        if "stack_top_b" in stack:
            stack["stack_top_b"] = scale_pairs(stack["stack_top_b"])

    if metric_key == "cpu":
        unit = metric_output["value_unit"]
        scale = unit_to_seconds(unit)
        if scale:
            scale_lists(metric_output, scale)
            metric_output["value_unit"] = "seconds"

    summary = {
        "meta": comparison.meta,
        "metric": metric_key,
        "metrics": {metric_key: metric_output},
    }

    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)
    else:
        print(json.dumps(summary, ensure_ascii=False, indent=2))