"""Deterministic synthetic graph generator for benchmarks.

Generates a social-network-like graph with 5:1 edge-to-node ratio.
All generation is seeded for reproducibility — every database gets
identical data for fair comparison.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

import numpy as np

from benchmarks.config import ScalePreset

# Node categories (roughly social-network-like)
CATEGORIES = ["Person", "Company", "Product", "Post", "Location"]
CATEGORY_WEIGHTS = [0.40, 0.10, 0.15, 0.20, 0.15]

# Edge types and their source → target category constraints
EDGE_TYPES = [
    ("KNOWS", "Person", "Person"),
    ("WORKS_AT", "Person", "Company"),
    ("PURCHASED", "Person", "Product"),
    ("AUTHORED", "Person", "Post"),
    ("LIVES_IN", "Person", "Location"),
    ("LOCATED_IN", "Company", "Location"),
    ("MENTIONS", "Post", "Product"),
    ("LIKES", "Person", "Post"),
]


@dataclass
class SyntheticData:
    """Container for generated benchmark data."""

    nodes: list[dict[str, Any]] = field(default_factory=list)
    edges: list[dict[str, Any]] = field(default_factory=list)
    vectors: np.ndarray | None = None
    scale: ScalePreset | None = None

    # Indexes for efficient lookup
    nodes_by_category: dict[str, list[int]] = field(default_factory=dict)


def generate(scale: ScalePreset) -> SyntheticData:
    """Generate a deterministic synthetic graph at the given scale.

    Returns a SyntheticData instance with nodes, edges, and vectors.
    """
    rng = np.random.default_rng(scale.seed)
    data = SyntheticData(scale=scale)

    # -- Generate nodes --------------------------------------------------
    categories = rng.choice(
        CATEGORIES,
        size=scale.num_nodes,
        p=CATEGORY_WEIGHTS,
    )

    for i in range(scale.num_nodes):
        cat = str(categories[i])
        node = {
            "bench_id": i,
            "category": cat,
            "name": f"{cat}_{i}",
            "score": float(rng.uniform(0, 100)),
            "active": bool(rng.random() > 0.2),
        }
        data.nodes.append(node)

        # Build category index
        if cat not in data.nodes_by_category:
            data.nodes_by_category[cat] = []
        data.nodes_by_category[cat].append(i)

    # -- Generate edges --------------------------------------------------
    # Build per-category node lists for constrained edge generation
    edges_generated = 0
    edge_set: set[tuple[int, int, str]] = set()

    while edges_generated < scale.num_edges:
        # Pick a random edge type
        et_idx = int(rng.integers(0, len(EDGE_TYPES)))
        rel_type, src_cat, tgt_cat = EDGE_TYPES[et_idx]

        src_pool = data.nodes_by_category.get(src_cat, [])
        tgt_pool = data.nodes_by_category.get(tgt_cat, [])

        if not src_pool or not tgt_pool:
            continue

        src_id = src_pool[int(rng.integers(0, len(src_pool)))]
        tgt_id = tgt_pool[int(rng.integers(0, len(tgt_pool)))]

        # No self-loops
        if src_id == tgt_id:
            continue

        # No duplicate edges of the same type
        edge_key = (src_id, tgt_id, rel_type)
        if edge_key in edge_set:
            continue

        edge_set.add(edge_key)
        data.edges.append({
            "source_id": src_id,
            "target_id": tgt_id,
            "rel_type": rel_type,
            "weight": float(rng.uniform(0.1, 1.0)),
            "timestamp": int(rng.integers(1_600_000_000, 1_700_000_000)),
        })
        edges_generated += 1

    # -- Generate vectors ------------------------------------------------
    # One vector per node, normalized for cosine similarity
    raw = rng.standard_normal((scale.num_nodes, scale.vec_dims)).astype(np.float32)
    norms = np.linalg.norm(raw, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    data.vectors = raw / norms

    return data


def generate_query_vector(scale: ScalePreset, query_seed: int = 99) -> np.ndarray:
    """Generate a deterministic query vector for vector search workloads."""
    rng = np.random.default_rng(query_seed)
    vec = rng.standard_normal(scale.vec_dims).astype(np.float32)
    vec /= np.linalg.norm(vec)
    return vec


def generate_workload_nodes(count: int, start_id: int, seed: int = 1000) -> list[dict[str, Any]]:
    """Generate a batch of nodes specifically for write workloads (W01)."""
    rng = np.random.default_rng(seed)
    nodes = []
    for i in range(count):
        nid = start_id + i
        nodes.append({
            "bench_id": nid,
            "category": "Person",
            "name": f"NewPerson_{nid}",
            "score": float(rng.uniform(0, 100)),
            "active": True,
        })
    return nodes


def generate_workload_edges(
    count: int,
    max_node_id: int,
    seed: int = 2000,
) -> list[dict[str, Any]]:
    """Generate a batch of edges specifically for write workloads (W02)."""
    rng = np.random.default_rng(seed)
    edges = []
    for _ in range(count):
        src = int(rng.integers(0, max_node_id))
        tgt = int(rng.integers(0, max_node_id))
        if src == tgt:
            tgt = (tgt + 1) % max_node_id
        edges.append({
            "source_id": src,
            "target_id": tgt,
            "rel_type": "KNOWS",
            "weight": float(rng.uniform(0.1, 1.0)),
            "timestamp": int(rng.integers(1_600_000_000, 1_700_000_000)),
        })
    return edges
