from tortoise import fields
from tortoise.models import Model
from .timestamp import Timestamp
from enum import Enum


class BotStatus(str, Enum):
    CONFIRMED = "confirmed"
    REJECTED = "rejected"
    PENDING = "pending"


class BotType(str, Enum):
    LLM = "llm"
    CHARACTER = "character"
    IMAGE = "image"
    VIDEO = "video"
    AUDIO = "audio"


class Bot(Model, Timestamp):
    id = fields.IntField(pk=True)
    name = fields.CharField(max_length=255, null=True)
    model = fields.CharField(max_length=255)
    style = fields.CharField(max_length=255, null=True)
    description = fields.TextField(null=True)
    image = fields.CharField(max_length=255, null=True)
    image_2 = fields.CharField(max_length=255, null=True)
    attachment = fields.IntField(default=0)
    attachment_type = fields.JSONField(null=True)
    tool = fields.BooleanField(default=False)
    public = fields.BooleanField(default=True)
    cost = fields.IntField(default=0)
    limit = fields.IntField(default=-1)
    order = fields.IntField(default=1)
    prompt = fields.TextField(null=True)
    links = fields.JSONField(null=True)
    docs = fields.JSONField(null=True)
    deleted = fields.BooleanField(default=False)
    web_search = fields.BooleanField(default=False)
    status = fields.CharEnumField(BotStatus, default=BotStatus.CONFIRMED)
    type = fields.CharEnumField(BotType, default=BotType.LLM)
    chatbots = fields.ReverseRelation["Chatbot"]
    marks = fields.ReverseRelation["Mark"]
    rates = fields.ReverseRelation["Rate"]
    allowed_levels = fields.JSONField(default=[1])
    banners = fields.ReverseRelation["Banner"]
    user = fields.ForeignKeyField(
        "models.User",
        related_name="bots",
        on_delete=fields.SET_NULL,
        null=True,
    )
    category = fields.ForeignKeyField(
        "models.Category",
        related_name="bots",
        on_delete=fields.SET_NULL,
        null=True,
    )

    class Meta:
        table = "bots"