import { Migration } from '../migrate.js'

export const up: Migration['up'] = async (db) => {
  // ML Models table
  await db.query(`
    CREATE TABLE IF NOT EXISTS ml_models (
      id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
      name VARCHAR(255) NOT NULL,
      version VARCHAR(100) NOT NULL,
      framework VARCHAR(100) NOT NULL,
      metadata JSONB DEFAULT '{}'::jsonb,
      artifact_path VARCHAR(500),
      created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
      updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
      UNIQUE(name, version)
    )
  `)

  // Model versions table
  await db.query(`
    CREATE TABLE IF NOT EXISTS model_versions (
      id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
      model_id UUID NOT NULL REFERENCES ml_models(id) ON DELETE CASCADE,
      version VARCHAR(100) NOT NULL,
      artifact_path VARCHAR(500),
      metadata JSONB DEFAULT '{}'::jsonb,
      created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
      UNIQUE(model_id, version)
    )
  `)

  // Model lineage table
  await db.query(`
    CREATE TABLE IF NOT EXISTS model_lineage (
      id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
      model_id UUID NOT NULL REFERENCES ml_models(id) ON DELETE CASCADE,
      training_job_id VARCHAR(255),
      parent_model_id UUID REFERENCES ml_models(id),
      metadata JSONB DEFAULT '{}'::jsonb,
      created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
    )
  `)

  // Indexes for ML models
  await db.query(`CREATE INDEX IF NOT EXISTS idx_ml_models_name ON ml_models(name)`)
  await db.query(`CREATE INDEX IF NOT EXISTS idx_model_versions_model ON model_versions(model_id)`)
  await db.query(`CREATE INDEX IF NOT EXISTS idx_model_lineage_model ON model_lineage(model_id)`)
}

export const down: Migration['down'] = async (db) => {
  await db.query(`DROP INDEX IF EXISTS idx_model_lineage_model`)
  await db.query(`DROP INDEX IF EXISTS idx_model_versions_model`)
  await db.query(`DROP INDEX IF EXISTS idx_ml_models_name`)
  await db.query(`DROP TABLE IF EXISTS model_lineage`)
  await db.query(`DROP TABLE IF EXISTS model_versions`)
  await db.query(`DROP TABLE IF EXISTS ml_models`)
}