from sqlalchemy import create_engine, Column, Integer, String, Float, Boolean, DateTime, ForeignKey
from sqlalchemy.orm import sessionmaker, relationship, declarative_base
from pydantic import BaseModel, ConfigDict
from typing import List
from datetime import datetime

# ایجاد مدل‌های SQLAlchemy
Base = declarative_base()

class Annotator(Base):
    __tablename__ = 'annotators'

    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False, unique=True)
    password = Column(String, nullable=False)
    last_login = Column(DateTime)
    is_active = Column(Boolean, default=True)
    annotations = relationship('Annotation', backref='annotator')

class AudioTrim(Base):
    __tablename__ = 'audio_trims'

    id = Column(Integer, primary_key=True)
    start = Column(Float, nullable=False)
    end = Column(Float, nullable=False)
    annotation_id = Column(Integer, ForeignKey('annotations.id'))

class Annotation(Base):
    __tablename__ = 'annotations'

    id = Column(Integer, primary_key=True)
    annotated_subtitle = Column(String, nullable=False)
    is_first_phase_accepted = Column(Boolean, nullable=False)
    create_at = Column(DateTime, nullable=False)
    update_at = Column(DateTime, nullable=False)
    audio_trims = relationship('AudioTrim', cascade='all, delete', backref='annotation')
    json_data_id = Column(Integer, ForeignKey('json_data.id'))
    annotator_id = Column(Integer, ForeignKey('annotators.id'))

class JsonData(Base):
    __tablename__ = 'json_data'

    id = Column(Integer, primary_key=True)
    voice_name = Column(String, nullable=False)
    original_subtitle = Column(String, nullable=False)
    ignore_it = Column(Boolean, nullable=False)
    is_approved_in_second_phase = Column(Boolean, nullable=False)
    annotations = relationship('Annotation', cascade='all, delete', backref='json_data')

# ایجاد مدل‌های Pydantic
class AudioTrimModel(BaseModel):
    start: float
    end: float
    model_config = ConfigDict(from_attributes=True)

class AnnotationModel(BaseModel):
    annotator: str
    annotated_subtitle: str
    is_first_phase_accepted: bool
    create_at: datetime
    update_at: datetime
    audio_trims: List[AudioTrimModel]
    model_config = ConfigDict(from_attributes=True)

class JsonDataModel(BaseModel):
    id: int
    voice_name: str
    original_subtitle: str
    ignore_it: bool
    is_approved_in_second_phase: bool
    annotations: List[AnnotationModel]
    model_config = ConfigDict(from_attributes=True)

# ایجاد جلسه
DB_URL = "mysql+pymysql://user:password@localhost/testdb"  # اطلاعات کاربری و دیتابیس را ویرایش کنید
engine = create_engine(DB_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def initialize_database():
    Base.metadata.create_all(bind=engine)

# تابع ذخیره داده‌ها در دیتابیس
def save_json_data(json_data_dict):
    db = SessionLocal()
    try:
        # استفاده از Pydantic برای اعتبارسنجی داده‌ها
        json_data = JsonDataModel(**json_data_dict)

        for annotation in json_data.annotations:
            # تامین موجودیت annotator
            db_annotator = db.query(Annotator).filter_by(name=annotation.annotator).first()
            if not db_annotator:
                db_annotator = Annotator(name=annotation.annotator, password="default_password")
                db.add(db_annotator)
                db.commit()  # اضافه کنید تا id اختصاص داده شود

            # تبدیل مدل Pydantic به مدل SQLAlchemy
            db_annotation = Annotation(
                annotated_subtitle=annotation.annotated_subtitle,
                is_first_phase_accepted=annotation.is_first_phase_accepted,
                create_at=annotation.create_at,
                update_at=annotation.update_at,
                audio_trims=[
                    AudioTrim(start=trim.start, end=trim.end) for trim in annotation.audio_trims
                ],
                annotator=db_annotator
            )

            db_json_data = JsonData(
                id=json_data.id,
                voice_name=json_data.voice_name,
                original_subtitle=json_data.original_subtitle,
                ignore_it=json_data.ignore_it,
                is_approved_in_second_phase=json_data.is_approved_in_second_phase,
                annotations=[db_annotation]
            )

            db.add(db_json_data)
            db.commit()

    except Exception as e:
        db.rollback()
        print(e)
    finally:
        db.close()

# فراخوانی تابع ایجاد جدول‌ها و افزودن داده
initialize_database()

# داده‌ی نمونه برای افزودن
json_data = {
    "id": 16562,
    "voice_name": "voice1.wav",
    "original_subtitle": "خلاصه با دلی ناراحت این کار رو کرد کاری رو که به نظرم میاد کار عقلانی بوده و رفتن رسیدن همون جایی که خاطره خداحافظی ۶ سال پیش رو واسه شکلتون زنده کرد اون بار به عنوان invalid اون بار به عنوان ادمی که از پس کار برنیومده",
    "ignore_it": False,
    "annotations": [
        {
            "annotator": "amin76",
            "annotated_subtitle": "خلاصه با دلی ناراحت این کار رو کرد کاری رو که به نظرم میاد کار عقلانی بوده و رفتن رسیدن همون جایی که خاطره خداحافظی شش سال پیش رو واسه شکلتون زنده کرد اون بار به عنوان اينوليد اون بار به عنوان آدمی که از پس کار برنیومده",
            "audio_trims": [
                {"start": 14.5, "end": 15.0},
                {"start": 14.3, "end": 15.0}
            ],
            "is_first_phase_accepted": False,
            "create_at": "2025-05-20T16:05:53.138811",
            "update_at": "2025-05-20T16:06:28.525393"
        }
    ],
    "is_approved_in_second_phase": False
}

save_json_data(json_data)