"""Shared fixtures for treeloom benchmarks."""

from __future__ import annotations

import pytest

from treeloom import CodePropertyGraph, CPGBuilder, NodeKind, TaintLabel, TaintPolicy
from treeloom.model.nodes import CpgNode


def _generate_python_source(num_classes: int, methods_per_class: int) -> bytes:
    """Generate deterministic synthetic Python source.

    Produces classes with methods that assign variables, call each other,
    branch on conditions, and return values — giving tree-sitter a realistic
    parse workload and treeloom a non-trivial graph to build.
    """
    lines: list[str] = [
        "# Synthetic benchmark source — generated by treeloom benchmarks",
        "",
        "import os",
        "import sys",
        "",
    ]

    # Emit a handful of standalone helper functions first so methods have
    # something external to call.
    for i in range(5):
        lines += [
            f"def helper_{i}(value, flag=False):",
            f"    result = value * {i + 1}",
            "    if flag:",
            f"        result = result + {i}",
            "    return result",
            "",
        ]

    for cls_idx in range(num_classes):
        lines += [
            f"class Worker{cls_idx}:",
            f'    """Worker class number {cls_idx}."""',
            "",
            "    def __init__(self, name, config):",
            "        self.name = name",
            "        self.config = config",
            f"        self._state = {cls_idx}",
            "",
        ]

        for meth_idx in range(methods_per_class):
            param_a = f"param_{meth_idx}_a"
            param_b = f"param_{meth_idx}_b"
            local_var = f"local_{meth_idx}"
            result_var = f"result_{meth_idx}"
            helper_idx = (cls_idx + meth_idx) % 5

            lines += [
                f"    def method_{meth_idx}(self, {param_a}, {param_b}=None):",
                f"        {local_var} = {param_a}",
                f"        if {param_b} is not None:",
                f"            {local_var} = helper_{helper_idx}({param_a}, flag=True)",
                f"        {result_var} = {local_var}",
                f"        for i in range({meth_idx + 1}):",
                f"            {result_var} = {result_var}",
                f"        return {result_var}",
                "",
            ]

        # Add a method that calls sibling methods so intra-class call edges form.
        lines += [
            "    def run(self, data):",
            "        output = data",
        ]
        for meth_idx in range(min(methods_per_class, 3)):
            lines.append(f"        output = self.method_{meth_idx}(output)")
        lines += [
            "        return output",
            "",
        ]

    return "\n".join(lines).encode()


def _build_cpg(source: bytes) -> CodePropertyGraph:
    return CPGBuilder().add_source(source, "synthetic.py", "python").build()


# ---------------------------------------------------------------------------
# Source fixtures at three sizes
# ---------------------------------------------------------------------------


@pytest.fixture(scope="session")
def source_small() -> bytes:
    """~500 LOC: 10 classes x 5 methods."""
    return _generate_python_source(num_classes=10, methods_per_class=5)


@pytest.fixture(scope="session")
def source_medium() -> bytes:
    """~2000 LOC: 30 classes x 8 methods."""
    return _generate_python_source(num_classes=30, methods_per_class=8)


@pytest.fixture(scope="session")
def source_large() -> bytes:
    """~5000 LOC: 70 classes x 10 methods."""
    return _generate_python_source(num_classes=70, methods_per_class=10)


# ---------------------------------------------------------------------------
# Pre-built CPG fixtures (session-scoped so each benchmark file shares them)
# ---------------------------------------------------------------------------


@pytest.fixture(scope="session")
def cpg_small(source_small: bytes) -> CodePropertyGraph:
    return _build_cpg(source_small)


@pytest.fixture(scope="session")
def cpg_medium(source_medium: bytes) -> CodePropertyGraph:
    return _build_cpg(source_medium)


# ---------------------------------------------------------------------------
# Taint policy fixture
# ---------------------------------------------------------------------------


@pytest.fixture(scope="session")
def taint_policy() -> TaintPolicy:
    """Simple policy: PARAMETER nodes are sources, CALL nodes are sinks."""

    def sources(node: CpgNode) -> TaintLabel | None:
        if node.kind == NodeKind.PARAMETER:
            return TaintLabel(name="param_taint", origin=node.id)
        return None

    def sinks(node: CpgNode) -> bool:
        return node.kind == NodeKind.CALL

    def sanitizers(node: CpgNode) -> bool:
        return False

    return TaintPolicy(sources=sources, sinks=sinks, sanitizers=sanitizers)
