from sqlalchemy import String, Integer, Boolean, text
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.dialects.postgresql import JSONB
from app.core.database import Base


class PlanLimits(Base):
    __tablename__ = "plan_limits"

    plan: Mapped[str] = mapped_column(String(50), primary_key=True)
    max_trees: Mapped[int | None] = mapped_column(Integer, nullable=True)
    max_sessions_per_month: Mapped[int | None] = mapped_column(Integer, nullable=True)
    max_users: Mapped[int | None] = mapped_column(Integer, nullable=True)
    custom_branding: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    priority_support: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    export_formats: Mapped[list] = mapped_column(JSONB, nullable=False, default=lambda: ["markdown", "text"])

    # AI Flow Builder limits
    max_ai_builds_per_month: Mapped[int | None] = mapped_column(Integer, nullable=True)
    max_ai_builds_per_24h: Mapped[int | None] = mapped_column(Integer, nullable=True)

    # KB Accelerator limits
    kb_accelerator_enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default=text("false"))
    kb_max_lifetime_conversions: Mapped[int | None] = mapped_column(Integer, nullable=True)
    kb_batch_max_size: Mapped[int | None] = mapped_column(Integer, nullable=True)
    kb_allowed_formats: Mapped[list] = mapped_column(JSONB, nullable=False, default=lambda: ["txt", "paste", "md"], server_default=text("'[\"txt\",\"paste\",\"md\"]'::jsonb"))
    kb_detailed_analysis: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default=text("false"))
    kb_conversational_refinement: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default=text("false"))
    kb_step_library_matching: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default=text("false"))
    kb_history_limit: Mapped[int | None] = mapped_column(Integer, nullable=True)