use sea_orm_migration::prelude::*;

pub struct Migration;

impl MigrationName for Migration {
    fn name(&self) -> &str {
        "ai_subagent_session"
    }
}

#[async_trait::async_trait]
impl MigrationTrait for Migration {
    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
        manager
            .get_connection()
            .execute_unprepared(
                r#"
create table if not exists ai_subagent_session
(
    id              uuid                     not null primary key,
    conversation_id uuid                     not null,
    message_id      uuid                     not null,
    children_id     varchar(255)             not null,
    role            varchar(64)              not null,
    task            text                     not null,
    output          text                     not null,
    input_tokens    bigint  default 0        not null,
    output_tokens   bigint  default 0        not null,
    model_name      varchar(255),
    status          varchar(32) default 'completed' not null,
    error_message   text,
    created_at      timestamp with time zone not null
);

create index if not exists idx_ai_subagent_session_conv
    on ai_subagent_session (conversation_id);

create index if not exists idx_ai_subagent_session_children
    on ai_subagent_session (children_id);

create index if not exists idx_ai_subagent_session_message
    on ai_subagent_session (message_id);
"#,
            )
            .await?;
        Ok(())
    }

    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
        manager
            .get_connection()
            .execute_unprepared(
                r#"
drop index if exists idx_ai_subagent_session_message;
drop index if exists idx_ai_subagent_session_children;
drop index if exists idx_ai_subagent_session_conv;
drop table if exists ai_subagent_session;
"#,
            )
            .await?;
        Ok(())
    }
}