"""Join Order Benchmark implementation.

This module provides the main benchmark class for the Join Order Benchmark,
which tests query optimizer join order selection using a complex schema based
on the Internet Movie Database (IMDB).

Copyright 2026 Joe Harris / BenchBox Project

Licensed under the MIT License. See LICENSE file in the project root for details.
"""

from __future__ import annotations

import time
from pathlib import Path
from typing import Any, Union

from benchbox.base import BaseBenchmark

from .generator import JoinOrderGenerator
from .queries import JoinOrderQueryManager
from .schema import JoinOrderSchema


class JoinOrderBenchmark(BaseBenchmark):
    """Join Order Benchmark implementation.

    The Join Order Benchmark is designed to test query optimizer
    join order selection capabilities using a complex schema with many
    interconnected tables. Originally based on the IMDB dataset, this
    implementation provides synthetic data generation while preserving
    the join patterns that stress-test query optimizers.
    """

    def __init__(
        self,
        scale_factor: float = 1.0,
        output_dir: Union[str, Path] | None = None,
        queries_dir: str | None = None,
        verbose: int | bool = 0,
        *,
        parallel: int = 1,
        force_regenerate: bool = False,
        **kwargs: Any,
    ) -> None:
        """Initialize the Join Order benchmark.

        Args:
            scale_factor: Scale factor for data generation (1.0 ≈ 1GB)
            output_dir: Directory for generated data files (defaults to benchmark_runs/datagen/joinorder_sf{X})
            queries_dir: Directory containing Join Order Benchmark query files (optional)
            verbose: Verbosity level (-v=1, -vv=2; bool True treated as 1)
            parallel: Number of parallel workers for data generation
            force_regenerate: Force regeneration even if data exists
            **kwargs: Additional options (compression settings, etc.) passed to generator
        """
        if not isinstance(parallel, int) or parallel < 1:
            raise ValueError(f"parallel must be a positive integer, got {parallel}")

        # Extract quiet from kwargs to avoid duplicate parameter
        quiet = kwargs.pop("quiet", False)

        super().__init__(
            scale_factor=scale_factor,
            output_dir=output_dir,
            verbose=verbose,
            quiet=quiet,
        )

        self.parallel = parallel
        self.force_regenerate = force_regenerate

        self.queries_dir = queries_dir
        self._schema = JoinOrderSchema()
        self._query_manager = JoinOrderQueryManager(queries_dir)

        # Pass all kwargs through to generator (includes compression params)
        generator_kwargs: dict[str, Any] = {
            "force_regenerate": force_regenerate,
            **kwargs,
        }

        self._generator = JoinOrderGenerator(
            scale_factor=scale_factor,
            output_dir=self.output_dir,
            verbose=verbose,
            quiet=quiet,
            **generator_kwargs,
        )

    def generate_data(self) -> list[Path]:
        """Generate Join Order Benchmark dataset.

        Returns:
            List of generated data file paths
        """
        self.log_verbose(f"Generating Join Order data at scale factor {self.scale_factor}...")

        start_time = time.time()
        data_files = self._generator.generate_data()
        generation_time = time.time() - start_time

        self.log_verbose(f"Generated {len(data_files)} data files in {generation_time:.2f}s")

        return data_files

    def get_schema(self, dialect: str = "sqlite") -> str:
        """Get database schema DDL.

        Args:
            dialect: Target SQL dialect

        Returns:
            DDL statements for creating all tables
        """
        return self._schema.get_create_tables_sql(dialect)

    def get_create_tables_sql(self, dialect: str = "sqlite") -> str:
        """Get CREATE TABLE statements for all tables.

        Args:
            dialect: Target SQL dialect

        Returns:
            SQL CREATE TABLE statements
        """
        return self._schema.get_create_tables_sql(dialect)

    def get_table_names(self) -> list[str]:
        """Get list of all table names.

        Returns:
            List of table names in the schema
        """
        return self._schema.get_table_names()

    def get_query(self, query_id: str, *, params: dict[str, Any] | None = None) -> str:
        """Get a specific query by ID.

        Args:
            query_id: Query identifier (e.g., '1a', '2b', etc.)
            params: Optional parameter values (not supported for JoinOrder)

        Returns:
            SQL query text

        Raises:
            ValueError: If params are provided
        """
        if params is not None:
            raise ValueError("JoinOrder queries are static and don't accept parameters")
        return self._query_manager.get_query(query_id)

    def get_queries(self) -> dict[str, str]:
        """Get all queries.

        Returns:
            Dictionary mapping query IDs to SQL text
        """
        return self._query_manager.get_all_queries()

    def get_query_ids(self) -> list[str]:
        """Get list of all query IDs.

        Returns:
            List of query IDs
        """
        return self._query_manager.get_query_ids()

    def get_query_count(self) -> int:
        """Get total number of queries.

        Returns:
            Number of queries available
        """
        return self._query_manager.get_query_count()

    def get_queries_by_complexity(self) -> dict[str, list[str]]:
        """Get queries categorized by complexity.

        Returns:
            Dictionary mapping complexity levels to query IDs
        """
        return self._query_manager.get_queries_by_complexity()

    def get_queries_by_pattern(self) -> dict[str, list[str]]:
        """Get queries categorized by join pattern.

        Returns:
            Dictionary mapping join patterns to query IDs
        """
        return self._query_manager.get_queries_by_pattern()

    def load_queries_from_directory(self, queries_dir: str) -> None:
        """Load queries from original Join Order Benchmark query files.

        Args:
            queries_dir: Path to directory containing Join Order Benchmark .sql files
        """
        self._query_manager = JoinOrderQueryManager(queries_dir)

    def get_table_info(self, table_name: str) -> dict[str, Any]:
        """Get information about a specific table.

        Args:
            table_name: Name of the table

        Returns:
            Dictionary with table schema information
        """
        return self._schema.get_table_info(table_name)

    def get_relationship_tables(self) -> list[str]:
        """Get list of relationship/junction tables.

        Returns:
            List of relationship table names
        """
        return self._schema.get_relationship_tables()

    def get_dimension_tables(self) -> list[str]:
        """Get list of main dimension tables.

        Returns:
            List of dimension table names
        """
        return self._schema.get_dimension_tables()

    def get_estimated_data_size(self) -> int:
        """Get estimated data size in bytes.

        Returns:
            Estimated total data size in bytes
        """
        return self._generator.get_total_size_estimate()

    def get_table_row_count(self, table_name: str) -> int:
        """Get expected row count for a table.

        Args:
            table_name: Name of the table

        Returns:
            Expected number of rows at current scale factor
        """
        return self._generator.get_table_row_count(table_name)

    def validate_query(self, query_id: str) -> bool:
        """Validate that a query is syntactically correct.

        Args:
            query_id: Query identifier

        Returns:
            True if query is valid, False otherwise
        """
        try:
            query = self.get_query(query_id)
            # Basic validation - check for SQL keywords
            query_upper = query.upper()
            required_keywords = ["SELECT", "FROM"]
            return all(keyword in query_upper for keyword in required_keywords)
        except Exception:
            return False

    def get_benchmark_info(self) -> dict[str, Any]:
        """Get comprehensive benchmark information.

        Returns:
            Dictionary with benchmark metadata
        """
        return {
            "benchmark_name": "Join Order Benchmark",
            "description": "Join order optimization benchmark using IMDB-like schema",
            "scale_factor": self.scale_factor,
            "output_dir": str(self.output_dir),
            "queries_dir": self.queries_dir,
            "total_queries": self.get_query_count(),
            "total_tables": len(self.get_table_names()),
            "relationship_tables": len(self.get_relationship_tables()),
            "dimension_tables": len(self.get_dimension_tables()),
            "estimated_size_bytes": self.get_estimated_data_size(),
            "query_complexity_distribution": self.get_queries_by_complexity(),
            "join_pattern_distribution": self.get_queries_by_pattern(),
            "reference_paper": "How Good Are Query Optimizers, Really? (VLDB 2015)",
            "authors": "Viktor Leis, Andrey Gubichev, Atanas Mirchev, Peter Boncz, Alfons Kemper, Thomas Neumann",
        }

    def __repr__(self) -> str:
        """String representation of the benchmark.

        Returns:
            String representation
        """
        return f"JoinOrderBenchmark(scale_factor={self.scale_factor}, queries={self.get_query_count()})"
