from sqlalchemy import JSON, Boolean, Column, DateTime, Index, Integer, String, Text
from sqlalchemy.dialects.mysql import MEDIUMBLOB
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.sql import func

Base = declarative_base()


class TestTask(Base):
    __tablename__ = "test_tasks"

    id = Column(Integer, primary_key=True, autoincrement=True)
    user_email = Column(String(255))
    cloud_provider = Column(String(50))
    region = Column(String(50))
    cluster_version = Column(String(50))
    org_id = Column(String(64))
    project_id = Column(String(64))

    tidb_spu = Column(String(50))
    tidb_count = Column(Integer)
    tidb_config = Column(JSON)

    tikv_spu = Column(String(50))
    tikv_count = Column(Integer)
    tikv_config = Column(JSON)
    tikv_storage_type = Column(String(50))
    tikv_storage_size_gi = Column(Integer)

    tiflash_spu = Column(String(50))
    tiflash_count = Column(Integer)
    tiflash_config = Column(JSON)
    tiflash_storage_type = Column(String(50))
    tiflash_storage_size_gi = Column(Integer)

    pd_spu = Column(String(50))
    pd_count = Column(Integer)
    pd_config = Column(JSON)

    tiproxy_spu = Column(String(50))
    tiproxy_count = Column(Integer)
    tiproxy_config = Column(JSON)

    workload_template = Column(String(50))
    workload_config = Column(JSON)
    workflow_execution_id = Column(String(100))
    workflow_execution_name = Column(String(255))
    status = Column(String(50), default="REQUESTED")
    deleted = Column(Boolean, nullable=False, default=False)

    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(
        DateTime(timezone=True), onupdate=func.now(), server_default=func.now()
    )

    # Indexes for common query patterns
    __table_args__ = (
        # For get_pending_tasks() and get_ended_tasks() - filter by status & deleted
        Index("idx_task_status_deleted", "status", "deleted"),
        # For get_all_tasks() - filter by deleted only
        Index("idx_task_deleted", "deleted"),
        # For report filtering - status, deleted, cloud, region, version
        Index(
            "idx_task_report_filter",
            "status",
            "deleted",
            "cloud_provider",
            "region",
            "cluster_version",
        ),
    )


class TestResult(Base):
    __tablename__ = "test_results"

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(
        Integer
    )  # Reference to test_tasks.id (no FK due to permission limits)
    result_type = Column(String(50))
    result_data = Column(JSON)
    error_msg = Column(Text)
    raw_execution_response = Column(JSON)
    deleted = Column(Boolean, nullable=False, default=False)
    notified = Column(Boolean, nullable=False, default=False)
    parsed = Column(Boolean, nullable=False, default=False)

    # Index for common query pattern: filter by task_id and deleted, order by id desc
    __table_args__ = (Index("idx_result_task_deleted_id", "task_id", "deleted", "id"),)


class ClusterMetrics(Base):
    __tablename__ = "cluster_metrics"

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(Integer, nullable=False)  # Reference to test_tasks.id
    cluster_id = Column(String(100), nullable=False)

    # Store compressed {diagram_name: {"dataframe": df, "queries": [...]}} dict
    # Use MEDIUMBLOB (16MB max) instead of BLOB (64KB max) for large metrics data
    # Made nullable to support S3 storage migration
    metrics_data = Column(MEDIUMBLOB, nullable=True)

    # S3 URI for storing large metrics data (e.g., s3://bucket/metrics/task_id/cluster_id.json.gz)
    # When set, metrics_data should be None
    metrics_s3_uri = Column(String(512), nullable=True)

    created_at = Column(DateTime(timezone=True), server_default=func.now())
    deleted = Column(Boolean, nullable=False, default=False)

    # Index for lookup by task_id + cluster_id
    # Business layer should ensure uniqueness of (task_id, cluster_id)
    __table_args__ = (
        Index("idx_cluster_metrics_task_cluster", "task_id", "cluster_id"),
        Index("idx_cluster_metrics_deleted", "deleted"),
    )