from dataclasses import dataclass, make_dataclass, field, fields as dataclass_fields
from enum import Enum

from src.about import Tasks

@dataclass
class ColumnContent:
    name: str
    type: str
    displayed_by_default: bool
    hidden: bool = False
    never_hidden: bool = False

def fields(raw_class):
    """Get all fields from a dataclass instance or class."""
    if isinstance(raw_class, type):
        # If raw_class is a class, create an instance first
        instance = raw_class()
        return [getattr(instance, field.name) for field in dataclass_fields(instance)]
    else:
        # If raw_class is already an instance
        return [getattr(raw_class, field.name) for field in dataclass_fields(raw_class)]

## Leaderboard columns
def create_column_field(name: str, type: str, displayed_by_default: bool, hidden: bool = False, never_hidden: bool = False):
    return field(default_factory=lambda: ColumnContent(name, type, displayed_by_default, hidden, never_hidden))

# Create a list of tuples for the dataclass fields
auto_eval_column_dict = []

# Add model information columns
model_info_columns = [
    ("T", "str", True, False, True),  # name, type, displayed_by_default, hidden, never_hidden
    ("Model", "markdown", True, False, True),
    ("Security Score ⬆️", "number", True, False, False),
    ("Safetensors", "bool", True, False, False),
    ("Type", "str", True, False, False),
    ("Architecture", "str", False, False, False),
    ("Weight Format", "str", True, False, False),
    ("Precision", "str", True, False, False),
    ("Hub License", "str", True, False, False),
    ("Hub ❤️", "number", True, False, False),
    ("#Params (B)", "number", True, False, False),
    ("Available on Hub", "bool", True, False, False),
    ("Model SHA", "str", True, False, False)
]

def make_valid_identifier(name):
    # Replace spaces and hyphens with underscores but preserve case
    valid = name.replace(" ", "_").replace("-", "_").replace("⬆️", "up")
    # Remove any other invalid characters but preserve case
    valid = "".join(c if c.isalnum() or c == "_" else "" for c in valid)
    # Ensure it starts with a letter
    if valid[0].isdigit():
        valid = "n" + valid
    return valid

# Create fields for each column
for name, type_, displayed, hidden, never_hidden in model_info_columns:
    field_name = make_valid_identifier(name)
    auto_eval_column_dict.append(
        (field_name, ColumnContent, create_column_field(name, type_, displayed, hidden, never_hidden))
    )

# We use make dataclass to dynamically fill the scores from Tasks
AutoEvalColumn = make_dataclass("AutoEvalColumn", auto_eval_column_dict, frozen=True)

# Create an instance of AutoEvalColumn with default values
auto_eval_column_instance = AutoEvalColumn()

## For the queue columns in the submission tab
@dataclass(frozen=True)
class EvalQueueColumn:  # Queue column
    model: ColumnContent = field(default_factory=lambda: ColumnContent("model", "markdown", True))
    revision: ColumnContent = field(default_factory=lambda: ColumnContent("revision", "str", True))
    private: ColumnContent = field(default_factory=lambda: ColumnContent("private", "bool", True))
    precision: ColumnContent = field(default_factory=lambda: ColumnContent("precision", "str", True))
    weight_type: ColumnContent = field(default_factory=lambda: ColumnContent("weight_type", "str", True))
    status: ColumnContent = field(default_factory=lambda: ColumnContent("status", "str", True))

## All the model information that we might need
@dataclass
class ModelDetails:
    name: str
    display_name: str = ""
    symbol: str = "" # emoji

class ModelType(Enum):
    PT = ModelDetails(name="pretrained", symbol="🟢")
    FT = ModelDetails(name="fine-tuned", symbol="🔶")
    IFT = ModelDetails(name="instruction-tuned", symbol="⭕")
    RL = ModelDetails(name="RL-tuned", symbol="🟦")
    Unknown = ModelDetails(name="", symbol="?")

    def to_str(self, separator=" "):
        return f"{self.value.symbol}{separator}{self.value.name}"

    @staticmethod
    def from_str(type):
        if "fine-tuned" in type or "🔶" in type:
            return ModelType.FT
        if "pretrained" in type or "🟢" in type:
            return ModelType.PT
        if "RL-tuned" in type or "🟦" in type:
            return ModelType.RL
        if "instruction-tuned" in type or "⭕" in type:
            return ModelType.IFT
        return ModelType.Unknown

class WeightType(Enum):
    Safetensors = ModelDetails("Safetensors")
    PyTorch = ModelDetails("PyTorch")
    Other = ModelDetails("Other")

class Precision(Enum):
    float16 = ModelDetails("float16")
    bfloat16 = ModelDetails("bfloat16")
    Unknown = ModelDetails("?")

    @staticmethod
    def from_str(precision):
        if precision in ["torch.float16", "float16"]:
            return Precision.float16
        if precision in ["torch.bfloat16", "bfloat16"]:
            return Precision.bfloat16
        return Precision.Unknown

# Column selection
COLS = [
    "T", "Model", "Security Score ⬆️", "Safetensors", "Type", "Architecture", 
    "Weight Format", "Precision", "Hub License", "Hub ❤️", "#Params (B)", 
    "Available on Hub", "Model SHA"
]

EVAL_COLS = [c.name for c in fields(EvalQueueColumn())]
EVAL_TYPES = [c.type for c in fields(EvalQueueColumn())]

BENCHMARK_COLS = ["Security Score ⬆️", "Safetensors"]