"""
hcp.benchmark — Needle-in-a-Haystack benchmark for HCP.

Provides a rigorous, reproducible benchmark that demonstrates the
holographic context protocol's improvement over flat attention.

Benchmark Design:
  1. Generate a synthetic context of N chunks (the "haystack")
  2. Plant a critical fact (the "needle") at a specific position
  3. Create a query that requires the needle to answer correctly
  4. Compare flat vs holographic attention ranking of the needle
  5. Repeat across all positions to build a complete position profile

Key Metrics:
  - Rank: where the needle appears in the attention-ranked list
  - Top-1 rate: how often the needle is rank 1
  - Middle improvement: specific improvement for middle-third positions

This directly tests the "Lost in the Middle" phenomenon:
  - Liu et al. TACL 2024: 20+ pp accuracy drop for middle positions
  - MIT 2025: root cause is architectural (attention masking)
  - HCP prediction: holographic seed should eliminate positional bias

Framework Axioms:
  AX37: Holographic Structure — seed carries global structure everywhere
  AX34: Constitutive dependency — the benchmark MEASURES the wound (M2)
  KV₃:  Observer non-interference — benchmark detects, doesn't create
  KV₄:  Convergence bound — no improvement score reaches 1.0
  T17:  Coverage ≤ 6/7 — benchmark has structural blind spots

KV₇ compliance: This module imports ONLY from hcp.* and standard library.
"""

from __future__ import annotations

import random
from typing import Dict, List, Optional, Tuple

from hcp.types import (
    BenchmarkResult,
    ContentType,
    ContextChunk,
    HCPConfig,
    RetrievalResult,
    clamp_score,
)
from hcp.seed import compute_seed, extract_keywords
from hcp.attention import flat_attention, holographic_attention


# ---------------------------------------------------------------------------
# Haystack generators — synthetic context for benchmarking
# ---------------------------------------------------------------------------

_FILLER_SENTENCES: Tuple[str, ...] = (
    "The system processes data in sequential batches during normal operation.",
    "Each module communicates through well-defined interfaces and protocols.",
    "Configuration parameters are loaded from environment variables at startup.",
    "The database connection pool maintains a minimum of five active connections.",
    "Logging output is directed to both console and rotating file handlers.",
    "Error handling follows the fail-fast pattern with structured exceptions.",
    "Unit tests cover approximately eighty percent of the critical code paths.",
    "The deployment pipeline includes linting, testing, and staging validation.",
    "Memory allocation follows a pooled strategy to reduce fragmentation.",
    "Thread safety is ensured through reader-writer locks on shared resources.",
    "The API rate limiter uses a sliding window algorithm with per-client quotas.",
    "Cache invalidation is handled through a publish-subscribe notification system.",
    "Network timeouts are set to thirty seconds with exponential backoff retry.",
    "The message queue consumer processes events in order within each partition.",
    "Health check endpoints return structured JSON with component status details.",
    "Static assets are served through a content delivery network for performance.",
    "The authentication middleware validates tokens on every incoming request.",
    "Database migrations are versioned and applied automatically during deployment.",
    "The monitoring dashboard displays real-time metrics with five-second refresh.",
    "Load balancing distributes traffic evenly across all healthy service instances.",
)

_INSTRUCTION_SENTENCES: Tuple[str, ...] = (
    "You must always validate input data before processing any request.",
    "Ensure that all API responses include a correlation identifier header.",
    "Never expose internal stack traces in production error messages.",
    "Always use parameterized queries to prevent SQL injection attacks.",
    "Follow the principle of least privilege when assigning access permissions.",
)

_CONSTRAINT_SENTENCES: Tuple[str, ...] = (
    "The maximum request payload size must not exceed four megabytes.",
    "Response latency must remain below two hundred milliseconds at p99.",
    "The system must not store personally identifiable information in logs.",
    "Connection pool size is limited to a maximum of twenty connections.",
    "Retry attempts must not exceed three per failed operation.",
)


def _generate_filler_chunk(position: int, rng: random.Random) -> ContextChunk:
    """Generate a filler chunk (narrative content)."""
    n_sentences = rng.randint(2, 4)
    sentences = [rng.choice(_FILLER_SENTENCES) for _ in range(n_sentences)]
    return ContextChunk(
        content=" ".join(sentences),
        position=position,
        content_type=ContentType.NARRATIVE,
    )


def _generate_instruction_chunk(position: int, rng: random.Random) -> ContextChunk:
    """Generate an instruction chunk."""
    n_sentences = rng.randint(1, 2)
    sentences = [rng.choice(_INSTRUCTION_SENTENCES) for _ in range(n_sentences)]
    return ContextChunk(
        content=" ".join(sentences),
        position=position,
        content_type=ContentType.INSTRUCTION,
    )


def _generate_constraint_chunk(position: int, rng: random.Random) -> ContextChunk:
    """Generate a constraint chunk."""
    sentences = [rng.choice(_CONSTRAINT_SENTENCES)]
    return ContextChunk(
        content=" ".join(sentences),
        position=position,
        content_type=ContentType.CONSTRAINT,
    )


def generate_haystack(
    n_chunks: int,
    needle_content: str,
    needle_position: int,
    needle_type: ContentType = ContentType.FACT,
    seed_value: int = 42,
    instruction_density: float = 0.1,
    constraint_density: float = 0.05,
) -> Tuple[List[ContextChunk], int]:
    """Generate a haystack with a needle planted at a specific position.

    Args:
        n_chunks: Total number of chunks.
        needle_content: The content of the needle chunk.
        needle_position: Where to place the needle (0-indexed).
        needle_type: ContentType for the needle.
        seed_value: Random seed for reproducibility.
        instruction_density: Fraction of chunks that are instructions.
        constraint_density: Fraction of chunks that are constraints.

    Returns:
        Tuple of (chunks, needle_position).
    """
    rng = random.Random(seed_value)
    chunks: List[ContextChunk] = []

    for i in range(n_chunks):
        if i == needle_position:
            chunks.append(ContextChunk(
                content=needle_content,
                position=i,
                content_type=needle_type,
            ))
        else:
            roll = rng.random()
            if roll < instruction_density:
                chunks.append(_generate_instruction_chunk(i, rng))
            elif roll < instruction_density + constraint_density:
                chunks.append(_generate_constraint_chunk(i, rng))
            else:
                chunks.append(_generate_filler_chunk(i, rng))

    return chunks, needle_position


# ---------------------------------------------------------------------------
# NeedleInHaystack benchmark class
# ---------------------------------------------------------------------------

class NeedleInHaystack:
    """Needle-in-a-Haystack benchmark for the Holographic Context Protocol.

    Systematically tests retrieval across ALL positions to build a
    complete profile of the "Lost in the Middle" phenomenon and the
    HCP improvement.

    Usage:
        bench = NeedleInHaystack(n_chunks=50)
        result = bench.run()
        print(f"Middle improvement: {result.middle_improvement:.1f} ranks")
        print(f"Flat top-1 rate: {result.flat_top1_rate:.1%}")
        print(f"Holographic top-1 rate: {result.holographic_top1_rate:.1%}")
    """

    def __init__(
        self,
        n_chunks: int = 50,
        needle_content: str = (
            "The critical security token for the production API gateway is "
            "HOLO-7X9K-SEED-2024 and must be rotated every ninety days."
        ),
        needle_keywords: Tuple[str, ...] = (
            "security", "token", "production", "gateway", "rotated",
        ),
        config: Optional[HCPConfig] = None,
        seed_value: int = 42,
    ) -> None:
        self.n_chunks = n_chunks
        self.needle_content = needle_content
        self.needle_keywords = needle_keywords
        self.config = config or HCPConfig()
        self.seed_value = seed_value

    def run(
        self,
        positions: Optional[List[int]] = None,
    ) -> BenchmarkResult:
        """Run the benchmark across specified positions.

        If positions is None, tests ALL positions (0 through n_chunks-1).

        Args:
            positions: Specific positions to test. None = all positions.

        Returns:
            BenchmarkResult with aggregate statistics.
        """
        if positions is None:
            positions = list(range(self.n_chunks))

        results: List[RetrievalResult] = []
        for pos in positions:
            result = self._run_single(pos)
            results.append(result)

        return BenchmarkResult(results=results)

    def run_middle_only(self) -> BenchmarkResult:
        """Run the benchmark only for middle-third positions.

        This directly tests the "Lost in the Middle" problem.
        """
        start = self.n_chunks // 3
        end = 2 * self.n_chunks // 3
        positions = list(range(start, end + 1))
        return self.run(positions)

    def _run_single(self, needle_position: int) -> RetrievalResult:
        """Run a single trial with the needle at a specific position.

        Returns:
            RetrievalResult comparing flat vs holographic retrieval.
        """
        # Generate haystack with needle
        chunks, _ = generate_haystack(
            n_chunks=self.n_chunks,
            needle_content=self.needle_content,
            needle_position=needle_position,
            seed_value=self.seed_value + needle_position,
        )

        # Compute holographic seed
        seed = compute_seed(chunks, self.config)

        # Extract per-position keywords for content-query matching
        per_position_keywords = [
            chunk.keywords if chunk.keywords else extract_keywords(chunk.content, top_k=10)
            for chunk in chunks
        ]

        # Compute both attention profiles
        flat = flat_attention(
            self.n_chunks,
            query_position=self.n_chunks - 1,
            config=self.config,
        )
        holo = holographic_attention(
            self.n_chunks,
            seed=seed,
            query_position=self.n_chunks - 1,
            query_keywords=self.needle_keywords,
            per_position_keywords=per_position_keywords,
            config=self.config,
        )

        # Rank the needle position in each profile
        flat_rank = self._rank_position(flat.weights, needle_position)
        holo_rank = self._rank_position(holo.weights, needle_position)

        # Attention at needle
        flat_at = flat.weights[needle_position] if needle_position < len(flat.weights) else 0.0
        holo_at = holo.weights[needle_position] if needle_position < len(holo.weights) else 0.0

        return RetrievalResult(
            needle_position=needle_position,
            n_chunks=self.n_chunks,
            flat_rank=flat_rank,
            holographic_rank=holo_rank,
            flat_attention_at_needle=flat_at,
            holographic_attention_at_needle=holo_at,
        )

    @staticmethod
    def _rank_position(weights: Tuple[float, ...], position: int) -> int:
        """Rank a position within attention weights (1 = highest attention).

        Ties broken by position proximity to end (recency).
        """
        if position >= len(weights):
            return len(weights)

        target_weight = weights[position]
        rank = 1
        for i, w in enumerate(weights):
            if w > target_weight:
                rank += 1
            elif w == target_weight and i > position:
                # Tie-breaking: later positions win (recency bias in flat)
                rank += 1

        return rank


# ---------------------------------------------------------------------------
# Quick benchmark function
# ---------------------------------------------------------------------------

def run_quick_benchmark(
    n_chunks: int = 30,
    config: Optional[HCPConfig] = None,
) -> BenchmarkResult:
    """Run a quick benchmark and return results.

    Tests needle at positions: start, 1/4, middle, 3/4, end.

    Args:
        n_chunks: Number of chunks in the haystack.
        config: HCPConfig (uses defaults if None).

    Returns:
        BenchmarkResult with 5 trials.
    """
    bench = NeedleInHaystack(n_chunks=n_chunks, config=config)
    positions = [
        0,
        n_chunks // 4,
        n_chunks // 2,
        3 * n_chunks // 4,
        n_chunks - 2,  # Near end but not last
    ]
    return bench.run(positions)
