#!/usr/bin/env python3
"""Compute pseudo-code token statistics using a HF tokenizer."""

from __future__ import annotations

import argparse
import json
import math
import statistics
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Sequence

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from transformers import AutoTokenizer


@dataclass(slots=True)
class PseudoSample:
    tokens: int
    pseudo_path: str | None
    function_name: str | None
    source_path: str | None
    origin: str


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Tokenize pseudo-code content from JSONL func-map files and emit "
            "rich statistics including textual, tabular, and graphical outputs."
        )
    )
    parser.add_argument(
        "-i",
        "--input",
        required=True,
        nargs="+",
        help="One or more func_map JSONL files containing pseudo content.",
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        type=Path,
        default=Path("pseudo_token_stats_output"),
        help="Directory where the text, table, and figure artefacts will be written.",
    )
    parser.add_argument(
        "--text-output",
        type=Path,
        default=None,
        help="Optional explicit path for the textual summary."
        " Defaults to OUTPUT_DIR/token_stats_summary.txt",
    )
    parser.add_argument(
        "--table-output",
        type=Path,
        default=None,
        help="Optional explicit path for the Markdown table summary."
        " Defaults to OUTPUT_DIR/token_stats_table.md",
    )
    parser.add_argument(
        "--figure-output",
        type=Path,
        default=None,
        help="Optional explicit path for the histogram figure."
        " Defaults to OUTPUT_DIR/token_distribution.png",
    )
    parser.add_argument(
        "--table-figure-output",
        type=Path,
        default=None,
        help=(
            "Optional explicit path for a high-quality table visualization (PNG)."
            " Defaults to OUTPUT_DIR/token_stats_table.png"
        ),
    )
    parser.add_argument(
        "--model",
        default="LLM4Binary/sk2decompile-struct-6.7b",
        help="Tokenizer identifier to use (HF repo or local path).",
    )
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=None,
        help="Optional huggingface cache directory override.",
    )
    parser.add_argument(
        "--local-files-only",
        action="store_true",
        help="Do not attempt to download the tokenizer, expect it to exist locally.",
    )
    parser.add_argument(
        "--bucket-size",
        type=int,
        default=128,
        help="Token bucket size used for histogram-friendly textual stats.",
    )
    parser.add_argument(
        "--max-top",
        type=int,
        default=10,
        help="How many of the longest pseudo entries to report in the summary.",
    )
    parser.add_argument(
        "--min-token",
        type=int,
        default=0,
        help="Skip pseudo entries tokenizing to <= this many tokens (default keeps all).",
    )
    return parser.parse_args()


def load_tokenizer(args: argparse.Namespace):
    kwargs = {"local_files_only": args.local_files_only}
    if args.cache_dir:
        kwargs["cache_dir"] = str(args.cache_dir)
    return AutoTokenizer.from_pretrained(args.model, **kwargs)


def read_jsonl(path: Path) -> Iterable[tuple[dict, int]]:
    with path.open("r", encoding="utf-8") as handle:
        for idx, raw in enumerate(handle, start=1):
            stripped = raw.strip()
            if not stripped:
                continue
            try:
                yield json.loads(stripped), idx
            except json.JSONDecodeError as exc:
                raise RuntimeError(f"Failed to parse {path}:{idx}: {exc}") from exc


def extract_pseudo_payload(entry: dict) -> tuple[str | None, str | None, str | None]:
    pseudo = entry.get("pseudo")
    if isinstance(pseudo, dict):
        return pseudo.get("content"), pseudo.get("path"), pseudo.get("function_name")
    if isinstance(pseudo, str):
        return pseudo, None, None
    return None, None, None


def percentile(data: Sequence[int], pct: float) -> float:
    if not data:
        return float("nan")
    if not 0 <= pct <= 100:
        raise ValueError("percentile must be between 0 and 100")
    if len(data) == 1:
        return float(data[0])
    position = (len(data) - 1) * pct / 100
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return float(data[int(position)])
    weight = position - lower
    return data[lower] * (1 - weight) + data[upper] * weight


def build_bucket_table(values: Sequence[int], bucket_size: int) -> list[tuple[str, int]]:
    counter = Counter()
    for val in values:
        bucket = val // bucket_size
        counter[bucket] += 1
    if not counter:
        return []
    max_bucket = max(counter)
    buckets: list[tuple[str, int]] = []
    for bucket in range(max_bucket + 1):
        start = bucket * bucket_size
        end = start + bucket_size - 1
        label = f"{start:>5}-{end:>5}"
        buckets.append((label, counter.get(bucket, 0)))
    return buckets


def render_text_summary(
    output_path: Path,
    model_name: str,
    samples: Sequence[PseudoSample],
    percentiles: list[tuple[int, float]],
    buckets: Sequence[tuple[str, int]],
    bucket_size: int,
    skipped: int,
    top_n: int,
) -> str:
    if samples:
        token_counts = [s.tokens for s in samples]
        mean_tokens = statistics.fmean(token_counts)
        median_tokens = statistics.median(token_counts)
        stdev_tokens = statistics.pstdev(token_counts)
        total_tokens = sum(token_counts)
        min_tokens = min(token_counts)
        max_tokens = max(token_counts)
    else:
        mean_tokens = median_tokens = stdev_tokens = 0.0
        total_tokens = 0
        min_tokens = max_tokens = 0

    lines: list[str] = []
    lines.append("Pseudo Token Statistics")
    lines.append("=" * 27)
    lines.append(f"Tokenizer          : {model_name}")
    lines.append(f"Samples processed  : {len(samples)}")
    lines.append(f"Skipped entries    : {skipped}")
    lines.append(f"Total tokens       : {total_tokens:,}")
    lines.append(f"Mean tokens/sample : {mean_tokens:.2f}")
    lines.append(f"Median tokens      : {median_tokens:.2f}")
    lines.append(f"Std deviation      : {stdev_tokens:.2f}")
    lines.append(f"Minimum tokens     : {min_tokens}")
    lines.append(f"Maximum tokens     : {max_tokens}")

    if percentiles:
        lines.append("")
        lines.append("Percentile breakdown (tokens)")
        lines.append("-" * 34)
        for pct, value in percentiles:
            lines.append(f"p{pct:>2} -> {value:.2f}")

    if buckets:
        lines.append("")
        lines.append(f"Bucket distribution (size={bucket_size})")
        lines.append("-" * 40)
        for label, count in buckets:
            lines.append(f"{label} : {count}")

    if samples and top_n > 0:
        top_samples = sorted(samples, key=lambda s: s.tokens, reverse=True)[:top_n]
        lines.append("")
        lines.append("Top pseudo entries by token count")
        lines.append("-" * 38)
        for idx, sample in enumerate(top_samples, start=1):
            descriptor = sample.function_name or "<unknown>"
            pseudo_path = sample.pseudo_path or "n/a"
            lines.append(
                f"{idx:>2}. {descriptor} | tokens={sample.tokens} | pseudo={pseudo_path}"
            )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return "\n".join(lines)


def render_table_markdown(
    output_path: Path,
    percentiles: Sequence[tuple[int, float]],
    buckets: Sequence[tuple[str, int]],
    bucket_size: int,
    total_samples: int,
) -> None:
    lines: list[str] = []
    lines.append("| Percentile | Token count |")
    lines.append("|-----------:|------------:|")
    for pct, value in percentiles:
        val = "n/a" if math.isnan(value) else f"{value:.2f}"
        lines.append(f"| p{pct:>2} | {val} |")

    lines.append("")
    lines.append(
        f"| Token buckets (size={bucket_size}) | Samples | Share (%) |"
    )
    lines.append("|:-------------------------------|--------:|---------:|")
    if buckets:
        for label, count in buckets:
            share = (count / total_samples * 100) if total_samples else 0
            lines.append(f"| {label} | {count} | {share:.2f} |")
    else:
        lines.append("| n/a | 0 | 0.00 |")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def render_table_figure(
    output_path: Path,
    percentiles: Sequence[tuple[int, float]],
    buckets: Sequence[tuple[str, int]],
    bucket_size: int,
    total_samples: int,
    model_name: str,
) -> None:
    if not percentiles and not buckets:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text("No data available for table rendering.\n", encoding="utf-8")
        return

    fig, axes = plt.subplots(2, 1, figsize=(11, 9), layout="constrained")
    headline_color = "#264653"
    accent_color = "#2a9d8f"

    for ax in axes:
        ax.axis("off")
        ax.set_facecolor("#f9fafb")

    pct_headers = [f"p{pct}" for pct, _ in percentiles]
    pct_values = ["n/a" if math.isnan(val) else f"{val:.2f}" for _, val in percentiles]
    if not pct_headers:
        pct_headers = ["Percentile"]
        pct_values = ["No data"]

    table_pct = axes[0].table(
        cellText=[pct_values],
        colLabels=pct_headers,
        loc="center",
        cellLoc="center",
    )
    table_pct.auto_set_font_size(False)
    table_pct.set_fontsize(11)
    table_pct.scale(1.0, 1.6)
    for (row, col), cell in table_pct.get_celld().items():
        if row == 0:
            cell.set_text_props(weight="bold", color="white")
            cell.set_facecolor(headline_color)
        else:
            cell.set_facecolor("#e9f5f2")
    axes[0].set_title(
        "Token count percentiles",
        loc="left",
        pad=10,
        fontsize=14,
        color=headline_color,
        fontweight="bold",
    )

    bucket_headers = ["Token range", "Samples", "Share (%)"]
    bucket_rows: list[list[str]] = []
    for label, count in buckets:
        share = (count / total_samples * 100) if total_samples else 0
        bucket_rows.append([label, f"{count}", f"{share:.2f}"])
    if not bucket_rows:
        bucket_rows = [["n/a", "0", "0.00"]]

    table_bucket = axes[1].table(
        cellText=bucket_rows,
        colLabels=bucket_headers,
        loc="center",
        cellLoc="center",
    )
    table_bucket.auto_set_font_size(False)
    table_bucket.set_fontsize(11)
    table_bucket.scale(1.0, 1.3)
    alt_colors = ["#fefae0", "#f4f1de"]
    for (row, col), cell in table_bucket.get_celld().items():
        if row == 0:
            cell.set_text_props(weight="bold", color="white")
            cell.set_facecolor(accent_color)
        else:
            cell.set_facecolor(alt_colors[(row - 1) % len(alt_colors)])

    axes[1].set_title(
        f"Bucket distribution (size={bucket_size})",
        loc="left",
        pad=10,
        fontsize=14,
        color=headline_color,
        fontweight="bold",
    )

    fig.suptitle(
        f"Pseudo token statistics – {model_name}",
        fontsize=12,
        color="#4a4a4a",
        fontweight="semibold",
    )
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=300)
    plt.close(fig)


def render_histogram(
    output_path: Path,
    token_counts: Sequence[int],
    model_name: str,
) -> None:
    if not token_counts:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text("No data to plot.\n", encoding="utf-8")
        return

    bins = min(75, max(10, int(math.sqrt(len(token_counts)))))
    fig, ax = plt.subplots(figsize=(10, 6), layout="constrained")
    colors = ["#264653", "#2a9d8f", "#e9c46a", "#f4a261", "#e76f51"]
    ax.hist(
        token_counts,
        bins=bins,
        color=colors[1],
        edgecolor="white",
        alpha=0.9,
    )
    mean_val = statistics.fmean(token_counts)
    median_val = statistics.median(token_counts)
    ax.axvline(mean_val, color=colors[3], linestyle="--", linewidth=1.2, label="mean")
    ax.axvline(median_val, color=colors[0], linestyle="-.", linewidth=1.2, label="median")
    ax.set_title("Pseudo token count distribution", fontsize=16, fontweight="bold")
    ax.set_xlabel("Token count", fontsize=12)
    ax.set_ylabel("Frequency", fontsize=12)
    ax.grid(alpha=0.3)
    ax.legend()
    ax.set_facecolor("#fdfdfd")
    fig.suptitle(f"Tokenizer: {model_name}", fontsize=11, color="#555555")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=300)
    plt.close(fig)


def run() -> None:
    args = parse_args()
    tokenizer = load_tokenizer(args)

    samples: List[PseudoSample] = []
    skipped = 0

    for input_path in args.input:
        path = Path(input_path)
        for record, line_no in read_jsonl(path):
            content, pseudo_path, pseudo_fn = extract_pseudo_payload(record)
            if not content:
                skipped += 1
                continue
            tokens = tokenizer(
                content,
                add_special_tokens=True,
                return_attention_mask=False,
                return_token_type_ids=False,
            )["input_ids"]
            token_count = len(tokens)
            if token_count <= args.min_token:
                skipped += 1
                continue
            source_path = None
            source = record.get("source")
            if isinstance(source, dict):
                source_path = source.get("path")
            samples.append(
                PseudoSample(
                    tokens=token_count,
                    pseudo_path=pseudo_path,
                    function_name=pseudo_fn,
                    source_path=source_path,
                    origin=f"{path}:{line_no}",
                )
            )

    sorted_samples = sorted(samples, key=lambda s: s.tokens)
    token_counts = [sample.tokens for sample in sorted_samples]
    pct_targets = [5, 10, 25, 50, 75, 90, 95, 99]
    pct_stats = [(pct, percentile(token_counts, pct)) for pct in pct_targets]
    bucket_table = build_bucket_table(token_counts, args.bucket_size)

    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    text_output = args.text_output or output_dir / "token_stats_summary.txt"
    table_output = args.table_output or output_dir / "token_stats_table.md"
    table_figure_output = (
        args.table_figure_output or output_dir / "token_stats_table.png"
    )
    figure_output = args.figure_output or output_dir / "token_distribution.png"

    summary = render_text_summary(
        text_output,
        args.model,
        sorted_samples,
        pct_stats,
        bucket_table,
        args.bucket_size,
        skipped,
        args.max_top,
    )
    print(summary)

    render_table_markdown(
        table_output,
        pct_stats,
        bucket_table,
        args.bucket_size,
        len(token_counts),
    )
    render_table_figure(
        table_figure_output,
        pct_stats,
        bucket_table,
        args.bucket_size,
        len(token_counts),
        args.model,
    )
    render_histogram(figure_output, token_counts, args.model)


if __name__ == "__main__":
    run()
