"""Simple benchmark table definitions for schema creation."""

import enum

from sqlalchemy import (
    JSON,
    Boolean,
    Column,
    DateTime,
    Enum,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    UniqueConstraint,
)
from sqlalchemy.sql import func


class BenchmarkStatus(enum.Enum):
    """Status of a benchmark run."""

    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    PAUSED = "paused"


class DatasetType(enum.Enum):
    """Supported dataset types."""

    SIMPLEQA = "simpleqa"
    BROWSECOMP = "browsecomp"
    CUSTOM = "custom"


# Simple table definitions for creation
benchmark_runs_table = {
    "table_name": "benchmark_runs",
    "columns": [
        Column("id", Integer, primary_key=True, index=True),
        Column("run_name", String(255), nullable=True),
        Column("config_hash", String(16), nullable=False, index=True),
        Column("query_hash_list", JSON, nullable=False),
        Column("search_config", JSON, nullable=False),
        Column("evaluation_config", JSON, nullable=False),
        Column("datasets_config", JSON, nullable=False),
        Column(
            "status",
            Enum(BenchmarkStatus),
            default=BenchmarkStatus.PENDING,
            nullable=False,
        ),
        Column(
            "created_at", DateTime, server_default=func.now(), nullable=False
        ),
        Column(
            "updated_at",
            DateTime,
            server_default=func.now(),
            onupdate=func.now(),
            nullable=False,
        ),
        Column("start_time", DateTime, nullable=True),
        Column("end_time", DateTime, nullable=True),
        Column("total_examples", Integer, default=0, nullable=False),
        Column("completed_examples", Integer, default=0, nullable=False),
        Column("failed_examples", Integer, default=0, nullable=False),
        Column("overall_accuracy", Float, nullable=True),
        Column("processing_rate", Float, nullable=True),
        Column("error_message", Text, nullable=True),
    ],
    "indexes": [
        Index("idx_benchmark_runs_config_hash", "config_hash"),
        Index("idx_benchmark_runs_status_created", "status", "created_at"),
    ],
}

benchmark_results_table = {
    "table_name": "benchmark_results",
    "columns": [
        Column("id", Integer, primary_key=True, index=True),
        Column(
            "benchmark_run_id",
            Integer,
            ForeignKey("benchmark_runs.id", ondelete="CASCADE"),
            nullable=False,
            index=True,
        ),
        Column("example_id", String(255), nullable=False),
        Column("query_hash", String(32), nullable=False, index=True),
        Column("dataset_type", Enum(DatasetType), nullable=False),
        Column("question", Text, nullable=False),
        Column("correct_answer", Text, nullable=False),
        Column("response", Text, nullable=True),
        Column("extracted_answer", Text, nullable=True),
        Column("confidence", String(10), nullable=True),
        Column("processing_time", Float, nullable=True),
        Column("sources", JSON, nullable=True),
        Column("is_correct", Boolean, nullable=True),
        Column("graded_confidence", String(10), nullable=True),
        Column("grader_response", Text, nullable=True),
        Column(
            "created_at", DateTime, server_default=func.now(), nullable=False
        ),
        Column("completed_at", DateTime, nullable=True),
        Column("research_error", Text, nullable=True),
        Column("evaluation_error", Text, nullable=True),
        Column("task_index", Integer, nullable=True),
        Column("result_metadata", JSON, nullable=True),
    ],
    "indexes": [
        Index(
            "idx_benchmark_results_run_dataset",
            "benchmark_run_id",
            "dataset_type",
        ),
        Index("idx_benchmark_results_query_hash", "query_hash"),
        Index("idx_benchmark_results_completed", "completed_at"),
    ],
    "constraints": [
        UniqueConstraint(
            "benchmark_run_id", "query_hash", name="uix_run_query"
        ),
    ],
}

benchmark_configs_table = {
    "table_name": "benchmark_configs",
    "columns": [
        Column("id", Integer, primary_key=True, index=True),
        Column("name", String(255), nullable=False),
        Column("description", Text, nullable=True),
        Column("config_hash", String(16), nullable=False, index=True),
        Column("search_config", JSON, nullable=False),
        Column("evaluation_config", JSON, nullable=False),
        Column("datasets_config", JSON, nullable=False),
        Column(
            "created_at", DateTime, server_default=func.now(), nullable=False
        ),
        Column(
            "updated_at",
            DateTime,
            server_default=func.now(),
            onupdate=func.now(),
            nullable=False,
        ),
        Column("is_default", Boolean, default=False, nullable=False),
        Column("is_public", Boolean, default=True, nullable=False),
        Column("usage_count", Integer, default=0, nullable=False),
        Column("last_used", DateTime, nullable=True),
        Column("best_accuracy", Float, nullable=True),
        Column("avg_processing_rate", Float, nullable=True),
    ],
    "indexes": [
        Index("idx_benchmark_configs_name", "name"),
        Index("idx_benchmark_configs_hash", "config_hash"),
        Index("idx_benchmark_configs_default", "is_default"),
    ],
}

benchmark_progress_table = {
    "table_name": "benchmark_progress",
    "columns": [
        Column("id", Integer, primary_key=True, index=True),
        Column(
            "benchmark_run_id",
            Integer,
            ForeignKey("benchmark_runs.id", ondelete="CASCADE"),
            nullable=False,
            index=True,
        ),
        Column(
            "timestamp", DateTime, server_default=func.now(), nullable=False
        ),
        Column("completed_examples", Integer, nullable=False),
        Column("total_examples", Integer, nullable=False),
        Column("overall_accuracy", Float, nullable=True),
        Column("dataset_accuracies", JSON, nullable=True),
        Column("processing_rate", Float, nullable=True),
        Column("estimated_completion", DateTime, nullable=True),
        Column("current_dataset", Enum(DatasetType), nullable=True),
        Column("current_example_id", String(255), nullable=True),
        Column("memory_usage", Float, nullable=True),
        Column("cpu_usage", Float, nullable=True),
    ],
    "indexes": [
        Index(
            "idx_benchmark_progress_run_time", "benchmark_run_id", "timestamp"
        ),
    ],
}


def create_benchmark_tables_simple(engine):
    """Create benchmark tables using simple table definitions."""
    from sqlalchemy import MetaData, Table

    metadata = MetaData()

    # Create tables
    tables_to_create = [
        benchmark_runs_table,
        benchmark_results_table,
        benchmark_configs_table,
        benchmark_progress_table,
    ]

    for table_def in tables_to_create:
        table = Table(
            table_def["table_name"],
            metadata,
            *table_def["columns"],
            extend_existing=True,
        )

        # Add indexes
        for index in table_def.get("indexes", []):
            index.table = table

        # Add constraints
        for constraint in table_def.get("constraints", []):
            table.append_constraint(constraint)

    # Create all tables
    metadata.create_all(engine, checkfirst=True)
