"""Storage experiment harness — benchmarks both backends on common queries.

Generates OHR datasets via the sampler and measures insert time, query time,
round-trip fidelity, and memory for both FlatComponentTable and JsonDictBackend.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any

from flexohr.core.ohr import OHR
from flexohr.core.sampler import OHRSampler, SamplerConfig
from flexohr.storage.flat_table import FlatComponentTable
from flexohr.storage.json_backend import JsonDictBackend

logger = logging.getLogger(__name__)


@dataclass
class BenchmarkResult:
    """Timing and memory results for a single benchmark run."""

    backend: str
    dataset_size: int
    insert_ms: float
    ohr_query_ms: float
    component_query_ms: float
    roundtrip_ms: float
    roundtrip_ok: bool
    memory_bytes: int
    component_count: int
    ohr_node_count: int

    def summary_dict(self) -> dict[str, Any]:
        return {
            "backend": self.backend,
            "n_ohrs": self.dataset_size,
            "n_components": self.component_count,
            "n_ohr_nodes": self.ohr_node_count,
            "insert_ms": round(self.insert_ms, 2),
            "ohr_query_ms": round(self.ohr_query_ms, 2),
            "comp_query_ms": round(self.component_query_ms, 2),
            "roundtrip_ms": round(self.roundtrip_ms, 2),
            "roundtrip_ok": self.roundtrip_ok,
            "memory_bytes": self.memory_bytes,
        }


class StorageExperiment:
    """Run storage benchmarks on both backends with sampler-generated OHRs."""

    def __init__(self, seed: int = 42) -> None:
        self.seed = seed

    def generate_dataset(self, n: int, config: SamplerConfig | None = None) -> list[OHR]:
        """Generate *n* OHRs using the sampler."""
        sampler = OHRSampler(config=config, seed=self.seed)
        return sampler.sample_batch(n)

    def benchmark_flat_table(self, ohrs: list[OHR]) -> BenchmarkResult:
        """Benchmark the flat component-per-row table backend."""
        # Insert
        t0 = time.perf_counter()
        table = FlatComponentTable()
        table.add_ohrs(ohrs)
        insert_ms = (time.perf_counter() - t0) * 1000

        # OHR-level query: find OHRs with a specific property value
        # Use a common chord quality as the target
        t0 = time.perf_counter()
        table.query_ohrs_by_property("chord_quality", "major_triad")
        ohr_query_ms = (time.perf_counter() - t0) * 1000

        # Component-level query: find all roots
        t0 = time.perf_counter()
        table.query_components_by_property("tone_function", "root")
        component_query_ms = (time.perf_counter() - t0) * 1000

        # Round-trip: reconstruct all top-level OHRs and compare dicts
        t0 = time.perf_counter()
        reconstructed = table.reconstruct_all_top_level()
        roundtrip_ms = (time.perf_counter() - t0) * 1000

        # Check round-trip by comparing dict representations
        roundtrip_ok = self._check_roundtrip(ohrs, reconstructed)

        return BenchmarkResult(
            backend="flat_table",
            dataset_size=len(ohrs),
            insert_ms=insert_ms,
            ohr_query_ms=ohr_query_ms,
            component_query_ms=component_query_ms,
            roundtrip_ms=roundtrip_ms,
            roundtrip_ok=roundtrip_ok,
            memory_bytes=table.memory_bytes,
            component_count=table.count_components(),
            ohr_node_count=table.count_ohrs(),
        )

    def benchmark_json_backend(self, ohrs: list[OHR]) -> BenchmarkResult:
        """Benchmark the JSON dict list backend."""
        # Insert
        t0 = time.perf_counter()
        backend = JsonDictBackend()
        backend.add_ohrs(ohrs)
        insert_ms = (time.perf_counter() - t0) * 1000

        # OHR-level query
        t0 = time.perf_counter()
        backend.query_ohrs_by_property("chord_quality", "major_triad")
        ohr_query_ms = (time.perf_counter() - t0) * 1000

        # Component-level query
        t0 = time.perf_counter()
        backend.query_components_by_property("tone_function", "root")
        component_query_ms = (time.perf_counter() - t0) * 1000

        # Round-trip
        t0 = time.perf_counter()
        reconstructed = backend.reconstruct_all()
        roundtrip_ms = (time.perf_counter() - t0) * 1000

        roundtrip_ok = self._check_roundtrip(ohrs, reconstructed)

        return BenchmarkResult(
            backend="json_dict",
            dataset_size=len(ohrs),
            insert_ms=insert_ms,
            ohr_query_ms=ohr_query_ms,
            component_query_ms=component_query_ms,
            roundtrip_ms=roundtrip_ms,
            roundtrip_ok=roundtrip_ok,
            memory_bytes=backend.memory_bytes,
            component_count=backend.count_components(),
            ohr_node_count=backend.count_ohrs(),
        )

    def run(
        self,
        sizes: list[int] | None = None,
        config: SamplerConfig | None = None,
    ) -> list[BenchmarkResult]:
        """Run full benchmark suite on both backends at multiple dataset sizes."""
        if sizes is None:
            sizes = [100, 1000]

        results: list[BenchmarkResult] = []
        for n in sizes:
            logger.info("Benchmarking with %d OHRs", n)
            ohrs = self.generate_dataset(n, config)

            results.append(self.benchmark_flat_table(ohrs))
            results.append(self.benchmark_json_backend(ohrs))

        return results

    def results_to_dataframe(self, results: list[BenchmarkResult]) -> Any:
        """Convert results to a pandas DataFrame for display."""
        import pandas as pd

        return pd.DataFrame([r.summary_dict() for r in results])

    # ── Internal ──

    def _check_roundtrip(self, originals: list[OHR], reconstructed: list[OHR]) -> bool:
        """Check round-trip fidelity by comparing dict representations."""
        if len(originals) != len(reconstructed):
            return False
        for orig, recon in zip(originals, reconstructed):
            if orig.to_dict() != recon.to_dict():
                return False
        return True
