from typing import Literal
from typing_extensions import Annotated
import sqlalchemy
from sqlalchemy.orm import mapped_column
from sqlalchemy import Integer, String, ForeignKey,func, DateTime, Boolean, LargeBinary
from sqlalchemy.orm import DeclarativeBase, Mapped
import uuid
import datetime
import pytz

# Set Jakarta timezone
jakarta_tz = pytz.timezone('Asia/Jakarta')

def get_jakarta_time():
    return datetime.datetime.now(jakarta_tz)

# Use the timezone-aware function in SQLAlchemy annotations
timestamp_current = Annotated[
    datetime.datetime,
    mapped_column(nullable=False, default=get_jakarta_time)  # Use default instead of server_default
]

timestamp_update = Annotated[
    datetime.datetime,
    mapped_column(nullable=False, default=get_jakarta_time, onupdate=get_jakarta_time)  # onupdate uses the Python function
]

message_role = Literal["user", "assistant"]

class Base(DeclarativeBase):
    type_annotation_map = {
        message_role: sqlalchemy.Enum("user", "assistant", name="message_role"),
    }

class User(Base):
    __tablename__ = "user"

    id = mapped_column(Integer, primary_key=True)
    name = mapped_column(String(100), nullable=False)
    username = mapped_column(String(100), unique=True, nullable=False)
    role_id = mapped_column(Integer, ForeignKey("role.id"))
    email = mapped_column(String(100), unique=True, nullable=False)
    password_hash = mapped_column(String(100), nullable=False)
    created_at: Mapped[timestamp_current] 
    updated_at : Mapped[timestamp_update] 

class Feedback(Base):
    __tablename__ = "feedback"
    
    id = mapped_column(Integer, primary_key=True)
    user_id = mapped_column(Integer, ForeignKey("user.id"))
    rating = mapped_column(Integer)
    comment = mapped_column(String(1000))
    created_at : Mapped[timestamp_current] 

class Role(Base):
    __tablename__ = "role"
    
    id = mapped_column(Integer, primary_key=True)
    role_name = mapped_column(String(200), nullable=False)
    description = mapped_column(String(200))
    
class User_Role(Base):
    __tablename__ = "user_role"
    
    id = mapped_column(Integer, primary_key=True)
    user_id = mapped_column(Integer, ForeignKey("user.id"))
    role_id = mapped_column(Integer, ForeignKey("role.id"))
    
class Bot(Base):
    __tablename__ = "bot"
    
    id = mapped_column(Integer, primary_key=True)
    user_id = mapped_column(Integer, ForeignKey("user.id"))
    bot_name = mapped_column(String(200), nullable=False)
    created_at : Mapped[timestamp_current] 
    
class Session(Base):
    __tablename__ = "session"
    
    id = mapped_column(String(36), primary_key=True, index=True, default=lambda: str(uuid.uuid4()))  # Store as string
    user_id = mapped_column(Integer, ForeignKey("user.id"))
    bot_id = mapped_column(Integer, ForeignKey("bot.id"))
    created_at : Mapped[timestamp_current]
    updated_at : Mapped[timestamp_update] 

class Message(Base):
     __tablename__ = "message"
     
     id = mapped_column(Integer, primary_key=True)
     session_id = mapped_column(String(36), ForeignKey("session.id"), nullable=False)  # Store as string
     role : Mapped[message_role]
     goal = mapped_column(String(200))
     created_at : Mapped[timestamp_current] 
    
class Category(Base):
    __tablename__ = "category"
    
    id = mapped_column(Integer, primary_key=True)
    category = mapped_column(String(200))
    created_at : Mapped[timestamp_current]
    updated_at : Mapped[timestamp_update]  

class Metadata(Base):
    __tablename__ = "metadata"
    
    id = mapped_column(Integer, primary_key=True)
    title = mapped_column(String(200))
    # image_data = mapped_column(LargeBinary, nullable=True)
    category_id = mapped_column(Integer, ForeignKey("category.id"))
    author = mapped_column(String(200))
    year = mapped_column(Integer)
    publisher = mapped_column(String(100))
    thumbnail = mapped_column(LargeBinary, nullable=True)
    created_at : Mapped[timestamp_current] 
    updated_at : Mapped[timestamp_update] 

     
class Bot_Meta(Base):
    __tablename__ = "bot_meta"
    
    id = mapped_column(Integer, primary_key=True)
    bot_id = mapped_column(Integer, ForeignKey("bot.id"))
    metadata_id = mapped_column(Integer, ForeignKey("metadata.id"))
    created_at : Mapped[timestamp_current]
    updated_at : Mapped[timestamp_update] 

class User_Meta(Base):
    __tablename__ = "user_meta"
    
    id = mapped_column(Integer, primary_key=True)
    user_id = mapped_column(Integer, ForeignKey("user.id"))
    metadata_id = mapped_column(Integer, ForeignKey("metadata.id"))
    created_at : Mapped[timestamp_current]
    updated_at : Mapped[timestamp_update] 
    
class Planning(Base):
    __tablename__="planning"
    
    id = mapped_column(Integer, primary_key=True)
    trials_id = mapped_column(Integer, ForeignKey("trials.id"))
    planning_name = mapped_column(String(200), nullable=False)
    duration = mapped_column(Integer, nullable=False)  # Duration in months
    start_date = mapped_column(DateTime, nullable=False)  # Start date of the planning
    end_date = mapped_column(DateTime, nullable=False)  # End date of the planning
    is_activated = mapped_column(Boolean)
    created_at : Mapped[timestamp_current]  # Automatically sets the current timestamp
    
class Trials(Base):
    __tablename__ = "trials"
    
    id = mapped_column(Integer, primary_key=True)
    token_used = mapped_column(Integer, nullable=False)  # Adjust length as needed
    token_planned = mapped_column(Integer, nullable=False)


class Session_Publisher(Base):
    __tablename__ = "session_publisher"
    
    id = mapped_column(String(36), primary_key=True, index=True, default=lambda: str(uuid.uuid4()))  # Store as string
    user_id = mapped_column(Integer, ForeignKey("user.id"))
    bot_name = mapped_column(String(100), nullable=True)
    metadata_id = mapped_column(Integer, ForeignKey("metadata.id"))
    created_at : Mapped[timestamp_current] 
    updated_at : Mapped[timestamp_update]