from sqlalchemy import Column, ForeignKey, Integer, String, Text
from sqlalchemy.dialects.postgresql import JSONB

from app.core.schema_binding import global_model
from app.models.base import BaseModel


@global_model
class TrainingDataset(BaseModel):
    __tablename__ = 'training_datasets'
    __table_args__ = {'schema': 'hub_global'}
    name = Column(String(255), nullable=False)
    description = Column(Text, nullable=True)
    dataset_type = Column(String(50), nullable=False)
    version = Column(Integer, default=1)
    parent_version_id = Column(Integer, nullable=True)
    total_samples = Column(Integer, default=0)
    train_samples = Column(Integer, default=0)
    validation_samples = Column(Integer, default=0)
    test_samples = Column(Integer, default=0)
    created_by = Column(Integer, nullable=True)
    meta_data = Column(JSONB, default={})

@global_model
class DatasetSample(BaseModel):
    __tablename__ = 'dataset_samples'
    __table_args__ = {'schema': 'hub_global'}
    dataset_id = Column(Integer, ForeignKey('hub_global.training_datasets.id'), nullable=False)
    content = Column(Text, nullable=False)
    label = Column(String(255), nullable=False)
    split = Column(String(50), default='train')
    source_task_id = Column(Integer, nullable=True)
    meta_data = Column(JSONB, default={})