import base64
import collections
import functools
import importlib.resources
import json
import linecache
import os
import sys
import sysconfig

from ._css_utils import get_combined_css
from .collector import Collector, extract_lineno
from .opcode_utils import get_opcode_mapping
from .string_table import StringTable


class StackTraceCollector(Collector):
    def __init__(self, sample_interval_usec, *, skip_idle=False):
        self.sample_interval_usec = sample_interval_usec
        self.skip_idle = skip_idle

    def collect(self, stack_frames, timestamps_us=None):
        weight = len(timestamps_us) if timestamps_us else 1
        for frames, thread_id in self._iter_stacks(stack_frames, skip_idle=self.skip_idle):
            self.process_frames(frames, thread_id, weight=weight)

    def process_frames(self, frames, thread_id, weight=1):
        pass


class CollapsedStackCollector(StackTraceCollector):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.stack_counter = collections.Counter()

    def process_frames(self, frames, thread_id, weight=1):
        # Extract only (filename, lineno, funcname) - opcode not needed for collapsed stacks
        # frame is (filename, location, funcname, opcode)
        call_tree = tuple(
            (f[0], extract_lineno(f[1]), f[2]) for f in reversed(frames)
        )
        self.stack_counter[(call_tree, thread_id)] += weight

    def export(self, filename):
        lines = []
        for (call_tree, thread_id), count in self.stack_counter.items():
            parts = [f"tid:{thread_id}"]
            for file, line, func in call_tree:
                # This is what pstats does for "special" frames:
                if file == "~" and line == 0:
                    part = func
                else:
                    part = f"{os.path.basename(file)}:{func}:{line}"
                parts.append(part)
            stack_str = ";".join(parts)
            lines.append((stack_str, count))

        lines.sort(key=lambda x: (-x[1], x[0]))

        with open(filename, "w") as f:
            for stack, count in lines:
                f.write(f"{stack} {count}\n")
        print(f"Collapsed stack output written to {filename}")


class FlamegraphCollector(StackTraceCollector):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.stats = {}
        self._root = {"samples": 0, "children": {}, "threads": set()}
        self._total_samples = 0
        self._sample_count = 0  # Track actual number of samples (not thread traces)
        self._func_intern = {}
        self._string_table = StringTable()
        self._all_threads = set()

        # Thread status statistics (similar to LiveStatsCollector)
        self.thread_status_counts = {
            "has_gil": 0,
            "on_cpu": 0,
            "gil_requested": 0,
            "unknown": 0,
            "has_exception": 0,
            "total": 0,
        }
        self.samples_with_gc_frames = 0

        # Per-thread statistics
        self.per_thread_stats = {}  # {thread_id: {has_gil, on_cpu, gil_requested, unknown, has_exception, total, gc_samples}}

    def collect(self, stack_frames, timestamps_us=None):
        """Override to track thread status statistics before processing frames."""
        # Weight is number of timestamps (samples with identical stack)
        weight = len(timestamps_us) if timestamps_us else 1

        # Increment sample count by weight
        self._sample_count += weight

        # Collect both aggregate and per-thread statistics using base method
        status_counts, has_gc_frame, per_thread_stats = self._collect_thread_status_stats(stack_frames)

        # Merge aggregate status counts (multiply by weight)
        for key in status_counts:
            self.thread_status_counts[key] += status_counts[key] * weight

        # Update aggregate GC frame count
        if has_gc_frame:
            self.samples_with_gc_frames += weight

        # Merge per-thread statistics (multiply by weight)
        for thread_id, stats in per_thread_stats.items():
            if thread_id not in self.per_thread_stats:
                self.per_thread_stats[thread_id] = {
                    "has_gil": 0,
                    "on_cpu": 0,
                    "gil_requested": 0,
                    "unknown": 0,
                    "has_exception": 0,
                    "total": 0,
                    "gc_samples": 0,
                }
            for key, value in stats.items():
                self.per_thread_stats[thread_id][key] += value * weight

        # Call parent collect to process frames
        super().collect(stack_frames, timestamps_us)

    def set_stats(self, sample_interval_usec, duration_sec, sample_rate,
                  error_rate=None, missed_samples=None, mode=None):
        """Set profiling statistics to include in flamegraph data."""
        self.stats = {
            "sample_interval_usec": sample_interval_usec,
            "duration_sec": duration_sec,
            "sample_rate": sample_rate,
            "error_rate": error_rate,
            "missed_samples": missed_samples,
            "mode": mode
        }

    def export(self, filename):
        flamegraph_data = self._convert_to_flamegraph_format()

        # Debug output with string table statistics
        num_functions = len(flamegraph_data.get("children", []))
        total_time = flamegraph_data.get("value", 0)
        string_count = len(self._string_table)
        s1 = "" if num_functions == 1 else "s"
        s2 = "" if total_time == 1 else "s"
        s3 = "" if string_count == 1 else "s"
        print(
            f"Flamegraph data: {num_functions} root function{s1}, "
            f"{total_time} total sample{s2}, "
            f"{string_count} unique string{s3}"
        )

        if num_functions == 0:
            print(
                "Warning: No functions found in profiling data. Check if sampling captured any data."
            )
            return

        html_content = self._create_flamegraph_html(flamegraph_data)

        with open(filename, "w", encoding="utf-8") as f:
            f.write(html_content)

        print(f"Flamegraph saved to: {filename}")

    @staticmethod
    @functools.lru_cache(maxsize=None)
    def _format_function_name(func):
        filename, lineno, funcname = func

        # Special frames like <GC> and <native> should not show file:line
        if filename == "~" and lineno == 0:
            return funcname

        if len(filename) > 50:
            parts = filename.split("/")
            if len(parts) > 2:
                filename = f".../{'/'.join(parts[-2:])}"

        return f"{funcname} ({filename}:{lineno})"

    def _convert_to_flamegraph_format(self):
        if self._total_samples == 0:
            return {
                "name": self._string_table.intern("No Data"),
                "value": 0,
                "children": [],
                "threads": [],
                "strings": self._string_table.get_strings()
            }

        def convert_children(children, min_samples):
            out = []
            for func, node in children.items():
                samples = node["samples"]
                if samples < min_samples:
                    continue

                # Intern all string components for maximum efficiency
                filename_idx = self._string_table.intern(func[0])
                funcname_idx = self._string_table.intern(func[2])
                name_idx = self._string_table.intern(self._format_function_name(func))

                child_entry = {
                    "name": name_idx,
                    "value": samples,
                    "self": node.get("self", 0),
                    "children": [],
                    "filename": filename_idx,
                    "lineno": func[1],
                    "funcname": funcname_idx,
                    "threads": sorted(list(node.get("threads", set()))),
                }

                source = self._get_source_lines(func)
                if source:
                    # Intern source lines for memory efficiency
                    source_indices = [self._string_table.intern(line) for line in source]
                    child_entry["source"] = source_indices

                # Include opcode data if available
                opcodes = node.get("opcodes", {})
                if opcodes:
                    child_entry["opcodes"] = dict(opcodes)

                # Recurse
                child_entry["children"] = convert_children(
                    node["children"], min_samples
                )
                out.append(child_entry)

            # Sort by value (descending) then by name index for consistent ordering
            out.sort(key=lambda x: (-x["value"], x["name"]))
            return out

        # Filter out very small functions (less than 0.1% of total samples)
        total_samples = self._total_samples
        min_samples = max(1, int(total_samples * 0.001))

        root_children = convert_children(self._root["children"], min_samples)
        if not root_children:
            return {
                "name": self._string_table.intern("No significant data"),
                "value": 0,
                "children": [],
                "strings": self._string_table.get_strings()
            }

        # Calculate thread status percentages for display
        is_free_threaded = bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
        total_threads = max(1, self.thread_status_counts["total"])
        thread_stats = {
            "has_gil_pct": (self.thread_status_counts["has_gil"] / total_threads) * 100,
            "on_cpu_pct": (self.thread_status_counts["on_cpu"] / total_threads) * 100,
            "gil_requested_pct": (self.thread_status_counts["gil_requested"] / total_threads) * 100,
            "has_exception_pct": (self.thread_status_counts["has_exception"] / total_threads) * 100,
            "gc_pct": (self.samples_with_gc_frames / max(1, self._sample_count)) * 100,
            "free_threaded": is_free_threaded,
            **self.thread_status_counts
        }

        # Calculate per-thread statistics with percentages
        per_thread_stats_with_pct = {}
        total_samples_denominator = max(1, self._sample_count)
        for thread_id, stats in self.per_thread_stats.items():
            total = max(1, stats["total"])
            per_thread_stats_with_pct[thread_id] = {
                "has_gil_pct": (stats["has_gil"] / total) * 100,
                "on_cpu_pct": (stats["on_cpu"] / total) * 100,
                "gil_requested_pct": (stats["gil_requested"] / total) * 100,
                "has_exception_pct": (stats["has_exception"] / total) * 100,
                "gc_pct": (stats["gc_samples"] / total_samples_denominator) * 100,
                **stats
            }

        # Build opcode mapping for JS
        opcode_mapping = get_opcode_mapping()

        # If we only have one root child, make it the root to avoid redundant level
        if len(root_children) == 1:
            main_child = root_children[0]
            # Update the name to indicate it's the program root
            old_name = self._string_table.get_string(main_child["name"])
            new_name = f"Program Root: {old_name}"
            main_child["name"] = self._string_table.intern(new_name)
            main_child["stats"] = {
                **self.stats,
                "thread_stats": thread_stats,
                "per_thread_stats": per_thread_stats_with_pct
            }
            main_child["threads"] = sorted(list(self._all_threads))
            main_child["strings"] = self._string_table.get_strings()
            main_child["opcode_mapping"] = opcode_mapping
            return main_child

        return {
            "name": self._string_table.intern("Program Root"),
            "value": total_samples,
            "children": root_children,
            "stats": {
                **self.stats,
                "thread_stats": thread_stats,
                "per_thread_stats": per_thread_stats_with_pct
            },
            "threads": sorted(list(self._all_threads)),
            "strings": self._string_table.get_strings(),
            "opcode_mapping": opcode_mapping
        }

    def process_frames(self, frames, thread_id, weight=1):
        """Process stack frames into flamegraph tree structure.

        Args:
            frames: List of (filename, location, funcname, opcode) tuples in
                    leaf-to-root order. location is (lineno, end_lineno, col_offset, end_col_offset).
                    opcode is None if not gathered.
            thread_id: Thread ID for this stack trace
            weight: Number of samples this stack represents (for batched RLE)
        """
        # Reverse to root->leaf order for tree building
        self._root["samples"] += weight
        self._total_samples += weight
        self._root["threads"].add(thread_id)
        self._all_threads.add(thread_id)

        current = self._root
        for filename, location, funcname, opcode in reversed(frames):
            lineno = extract_lineno(location)
            func = (filename, lineno, funcname)
            func = self._func_intern.setdefault(func, func)

            node = current["children"].get(func)
            if node is None:
                node = {"samples": 0, "children": {}, "threads": set(), "opcodes": collections.Counter(), "self": 0}
                current["children"][func] = node
            node["samples"] += weight
            node["threads"].add(thread_id)

            if opcode is not None:
                node["opcodes"][opcode] += weight

            current = node

        if current is not self._root:
            current["self"] += weight

    def _get_source_lines(self, func):
        filename, lineno, _ = func

        try:
            lines = []
            start_line = max(1, lineno - 2)
            end_line = lineno + 3

            for line_num in range(start_line, end_line):
                line = linecache.getline(filename, line_num)
                if line.strip():
                    marker = "→ " if line_num == lineno else "  "
                    lines.append(f"{marker}{line_num}: {line.rstrip()}")

            return lines if lines else None

        except Exception:
            return None

    def _create_flamegraph_html(self, data):
        data_json = json.dumps(data)

        template_dir = importlib.resources.files(__package__)
        vendor_dir = template_dir / "_vendor"
        assets_dir = template_dir / "_assets"

        d3_path = vendor_dir / "d3" / "7.8.5" / "d3.min.js"
        d3_flame_graph_dir = vendor_dir /  "d3-flame-graph" / "4.1.3"
        fg_css_path = d3_flame_graph_dir / "d3-flamegraph.css"
        fg_js_path = d3_flame_graph_dir / "d3-flamegraph.min.js"
        fg_tooltip_js_path = d3_flame_graph_dir / "d3-flamegraph-tooltip.min.js"

        html_template = (template_dir / "_flamegraph_assets" / "flamegraph_template.html").read_text(encoding="utf-8")
        css_content = get_combined_css("flamegraph")
        base_js = (template_dir / "_shared_assets" / "base.js").read_text(encoding="utf-8")
        component_js = (template_dir / "_flamegraph_assets" / "flamegraph.js").read_text(encoding="utf-8")
        js_content = f"{base_js}\n{component_js}"

        # Set title and subtitle based on whether this is a differential flamegraph
        is_differential = data.get("stats", {}).get("is_differential", False)
        if is_differential:
            title = "Tachyon Profiler - Differential Flamegraph Report"
            subtitle = "Differential Flamegraph Report"
        else:
            title = "Tachyon Profiler - Flamegraph Report"
            subtitle = "Flamegraph Report"

        html_template = html_template.replace("{{TITLE}}", title)
        html_template = html_template.replace("{{SUBTITLE}}", subtitle)

        # Inline first-party CSS/JS
        html_template = html_template.replace(
            "<!-- INLINE_CSS -->", f"<style>\n{css_content}\n</style>"
        )
        html_template = html_template.replace(
            "<!-- INLINE_JS -->", f"<script>\n{js_content}\n</script>"
        )

        png_path = assets_dir / "tachyon-logo.png"
        b64_logo = base64.b64encode(png_path.read_bytes()).decode("ascii")

        # Let CSS control size; keep markup simple
        logo_html = f'<img src="data:image/png;base64,{b64_logo}" alt="Tachyon logo"/>'
        html_template = html_template.replace("<!-- INLINE_LOGO -->", logo_html)
        html_template = html_template.replace(
            "<!-- PYTHON_VERSION -->", f"{sys.version_info.major}.{sys.version_info.minor}"
        )

        d3_js = d3_path.read_text(encoding="utf-8")
        fg_css = fg_css_path.read_text(encoding="utf-8")
        fg_js = fg_js_path.read_text(encoding="utf-8")
        fg_tooltip_js = fg_tooltip_js_path.read_text(encoding="utf-8")

        html_template = html_template.replace(
            "<!-- INLINE_VENDOR_D3_JS -->",
            f"<script>\n{d3_js}\n</script>",
        )
        html_template = html_template.replace(
            "<!-- INLINE_VENDOR_FLAMEGRAPH_CSS -->",
            f"<style>\n{fg_css}\n</style>",
        )
        html_template = html_template.replace(
            "<!-- INLINE_VENDOR_FLAMEGRAPH_JS -->",
            f"<script>\n{fg_js}\n</script>",
        )
        html_template = html_template.replace(
            "<!-- INLINE_VENDOR_FLAMEGRAPH_TOOLTIP_JS -->",
            f"<script>\n{fg_tooltip_js}\n</script>",
        )

        # Replace the placeholder with actual data
        html_content = html_template.replace(
            "{{FLAMEGRAPH_DATA}}", data_json
        )

        return html_content


class DiffFlamegraphCollector(FlamegraphCollector):
    """Differential flamegraph collector that compares against a baseline binary profile."""

    def __init__(self, sample_interval_usec, *, baseline_binary_path, skip_idle=False):
        super().__init__(sample_interval_usec, skip_idle=skip_idle)
        if not os.path.exists(baseline_binary_path):
            raise ValueError(f"Baseline file not found: {baseline_binary_path}")
        self.baseline_binary_path = baseline_binary_path
        self._baseline_collector = None
        self._elided_paths = set()

    def _load_baseline(self):
        """Load baseline profile from binary file."""
        from .binary_reader import BinaryReader

        with BinaryReader(self.baseline_binary_path) as reader:
            info = reader.get_info()

            baseline_collector = FlamegraphCollector(
                sample_interval_usec=info['sample_interval_us'],
                skip_idle=self.skip_idle
            )

            reader.replay_samples(baseline_collector)

        self._baseline_collector = baseline_collector

    def _aggregate_path_samples(self, root_node, path=None):
        """Aggregate samples by stack path, excluding line numbers for cross-profile matching."""
        if path is None:
            path = ()

        stats = {}

        for func, node in root_node["children"].items():
            filename, _lineno, funcname = func
            func_key = (filename, funcname)
            path_key = path + (func_key,)

            total_samples = node.get("samples", 0)
            self_samples = node.get("self", 0)

            if path_key in stats:
                stats[path_key]["total"] += total_samples
                stats[path_key]["self"] += self_samples
            else:
                stats[path_key] = {
                    "total": total_samples,
                    "self": self_samples
                }

            child_stats = self._aggregate_path_samples(node, path_key)
            for key, data in child_stats.items():
                if key in stats:
                    stats[key]["total"] += data["total"]
                    stats[key]["self"] += data["self"]
                else:
                    stats[key] = data

        return stats

    def _convert_to_flamegraph_format(self):
        """Convert to flamegraph format with differential annotations."""
        if self._baseline_collector is None:
            self._load_baseline()

        current_flamegraph = super()._convert_to_flamegraph_format()

        current_stats = self._aggregate_path_samples(self._root)
        baseline_stats = self._aggregate_path_samples(self._baseline_collector._root)

        # Scale baseline values to make them comparable, accounting for both
        # sample count differences and sample interval differences.
        baseline_total = self._baseline_collector._total_samples
        if baseline_total > 0 and self._total_samples > 0:
            current_time = self._total_samples * self.sample_interval_usec
            baseline_time = baseline_total * self._baseline_collector.sample_interval_usec
            scale = current_time / baseline_time
        elif baseline_total > 0:
            # Current profile is empty - use interval-based scale for elided display
            scale = self.sample_interval_usec / self._baseline_collector.sample_interval_usec
        else:
            scale = 1.0

        self._annotate_nodes_with_diff(current_flamegraph, current_stats, baseline_stats, scale)
        self._add_elided_flamegraph(current_flamegraph, current_stats, baseline_stats, scale)

        return current_flamegraph

    def _annotate_nodes_with_diff(self, current_flamegraph, current_stats, baseline_stats, scale):
        """Annotate each node in the tree with diff metadata."""
        if "stats" not in current_flamegraph:
            current_flamegraph["stats"] = {}

        current_flamegraph["stats"]["baseline_samples"] = self._baseline_collector._total_samples
        current_flamegraph["stats"]["current_samples"] = self._total_samples
        current_flamegraph["stats"]["baseline_scale"] = scale
        current_flamegraph["stats"]["is_differential"] = True

        if self._is_promoted_root(current_flamegraph):
            self._add_diff_data_to_node(current_flamegraph, (), current_stats, baseline_stats, scale)
        else:
            for child in current_flamegraph["children"]:
                self._add_diff_data_to_node(child, (), current_stats, baseline_stats, scale)

    def _add_diff_data_to_node(self, node, path, current_stats, baseline_stats, scale):
        """Recursively add diff metadata to nodes."""
        func_key = self._extract_func_key(node, self._string_table)
        path_key = path + (func_key,) if func_key else path

        current_data = current_stats.get(path_key, {"total": 0, "self": 0})
        baseline_data = baseline_stats.get(path_key, {"total": 0, "self": 0})

        current_self = current_data["self"]
        baseline_self = baseline_data["self"] * scale
        baseline_total = baseline_data["total"] * scale

        diff = current_self - baseline_self
        if baseline_self > 0:
            diff_pct = (diff / baseline_self) * 100.0
        elif current_self > 0:
            diff_pct = 100.0
        else:
            diff_pct = 0.0

        node["baseline"] = baseline_self
        node["baseline_total"] = baseline_total
        node["self_time"] = current_self
        node["diff"] = diff
        node["diff_pct"] = diff_pct

        if "children" in node and node["children"]:
            for child in node["children"]:
                self._add_diff_data_to_node(child, path_key, current_stats, baseline_stats, scale)

    def _is_promoted_root(self, data):
        """Check if the data represents a promoted root node."""
        return "filename" in data and "funcname" in data

    def _add_elided_flamegraph(self, current_flamegraph, current_stats, baseline_stats, scale):
        """Calculate elided paths and add elided flamegraph to stats."""
        self._elided_paths = baseline_stats.keys() - current_stats.keys()

        current_flamegraph["stats"]["elided_count"] = len(self._elided_paths)

        if self._elided_paths:
            elided_flamegraph = self._build_elided_flamegraph(baseline_stats, scale)
            if elided_flamegraph:
                current_flamegraph["stats"]["elided_flamegraph"] = elided_flamegraph

    def _build_elided_flamegraph(self, baseline_stats, scale):
        """Build flamegraph containing only elided paths from baseline.

        This re-runs the base conversion pipeline on the baseline collector
        to produce a complete formatted flamegraph, then prunes it to keep
        only elided paths.
        """
        if not self._baseline_collector or not self._elided_paths:
            return None

        # Suppress source line collection for elided nodes - these functions
        # no longer exist in the current profile, so source lines from the
        # current machine's filesystem would be misleading or unavailable.
        orig_get_source = self._baseline_collector._get_source_lines
        self._baseline_collector._get_source_lines = lambda func: None
        try:
            baseline_data = self._baseline_collector._convert_to_flamegraph_format()
        finally:
            self._baseline_collector._get_source_lines = orig_get_source

        # Remove non-elided nodes and recalculate values
        if not self._extract_elided_nodes(baseline_data, path=()):
            return None

        self._add_elided_metadata(baseline_data, baseline_stats, scale, path=())

        # Merge only profiling metadata, not thread-level stats
        for key in ("sample_interval_usec", "duration_sec", "sample_rate",
                    "error_rate", "missed_samples", "mode"):
            if key in self.stats:
                baseline_data["stats"][key] = self.stats[key]
        baseline_data["stats"]["is_differential"] = True
        baseline_data["stats"]["baseline_samples"] = self._baseline_collector._total_samples
        baseline_data["stats"]["current_samples"] = self._total_samples

        return baseline_data

    def _extract_elided_nodes(self, node, path):
        """Remove non-elided nodes and recalculate values bottom-up."""
        if not node:
            return False

        func_key = self._extract_func_key(node, self._baseline_collector._string_table)
        current_path = path + (func_key,) if func_key else path

        is_elided = current_path in self._elided_paths if func_key else False

        if "children" in node:
            # Filter children, keeping only those with elided descendants
            elided_children = []
            total_value = 0
            for child in node["children"]:
                if self._extract_elided_nodes(child, current_path):
                    elided_children.append(child)
                    total_value += child.get("value", 0)
            node["children"] = elided_children

            # Recalculate value for structural (non-elided) ancestor nodes;
            # elided nodes keep their original value to preserve self-samples
            if elided_children and not is_elided:
                node["value"] = total_value

        # Keep this node if it's elided or has elided descendants
        return is_elided or bool(node.get("children"))

    def _add_elided_metadata(self, node, baseline_stats, scale, path):
        """Add differential metadata showing this path disappeared."""
        if not node:
            return

        func_key = self._extract_func_key(node, self._baseline_collector._string_table)
        current_path = path + (func_key,) if func_key else path

        if func_key and current_path in baseline_stats:
            baseline_data = baseline_stats[current_path]
            baseline_self = baseline_data["self"] * scale
            baseline_total = baseline_data["total"] * scale

            node["baseline"] = baseline_self
            node["baseline_total"] = baseline_total
            node["diff"] = -baseline_self
        else:
            node["baseline"] = 0
            node["baseline_total"] = 0
            node["diff"] = 0

        node["self_time"] = 0
        # Elided paths have zero current self-time, so the change is always
        # -100% when there was actual baseline self-time to lose.
        # For internal nodes with no baseline self-time, use 0% to avoid
        # misleading tooltips.
        if baseline_self > 0:
            node["diff_pct"] = -100.0
        else:
            node["diff_pct"] = 0.0

        if "children" in node and node["children"]:
            for child in node["children"]:
                self._add_elided_metadata(child, baseline_stats, scale, current_path)

    def _extract_func_key(self, node, string_table):
        """Extract (filename, funcname) key from node, excluding line numbers.

        Line numbers are excluded to match functions even if they moved.
        Returns None for root nodes that don't have function information.
        """
        if "filename" not in node or "funcname" not in node:
            return None
        filename = string_table.get_string(node["filename"])
        funcname = string_table.get_string(node["funcname"])
        return (filename, funcname)