import { Context } from '../types/context'

export interface MLModel {
  id: string
  name: string
  version: string
  framework: string
  metadata: Record<string, any>
  artifactPath?: string
  createdAt: Date
  updatedAt: Date
}

export interface ModelVersion {
  id: string
  modelId: string
  version: string
  artifactPath?: string
  metadata: Record<string, any>
  createdAt: Date
}

export interface CreateModelInput {
  name: string
  version: string
  framework: string
  metadata?: Record<string, any>
  artifactPath?: string
}

export async function getModels(context: Context, name?: string): Promise<MLModel[]> {
  const db = context.db
  let query = 'SELECT * FROM ml_models WHERE 1=1'
  const params: any[] = []

  if (name) {
    query += ' AND name = $1'
    params.push(name)
  }

  query += ' ORDER BY created_at DESC'

  const result = await db.query(query, params)
  return result.rows.map(mapModel)
}

export async function getModel(context: Context, id: string): Promise<MLModel> {
  const db = context.db
  const result = await db.query('SELECT * FROM ml_models WHERE id = $1', [id])

  if (result.rows.length === 0) {
    throw new Error('Model not found')
  }

  return mapModel(result.rows[0])
}

export async function createModel(context: Context, input: CreateModelInput): Promise<MLModel> {
  const db = context.db
  const result = await db.query(
    `INSERT INTO ml_models (name, version, framework, metadata, artifact_path)
     VALUES ($1, $2, $3, $4, $5)
     RETURNING *`,
    [
      input.name,
      input.version,
      input.framework,
      JSON.stringify(input.metadata || {}),
      input.artifactPath || null,
    ]
  )

  return mapModel(result.rows[0])
}

export async function getModelVersions(context: Context, modelId: string): Promise<ModelVersion[]> {
  const db = context.db
  const result = await db.query(
    'SELECT * FROM model_versions WHERE model_id = $1 ORDER BY created_at DESC',
    [modelId]
  )
  return result.rows.map(mapModelVersion)
}

function mapModel(row: any): MLModel {
  return {
    id: row.id,
    name: row.name,
    version: row.version,
    framework: row.framework,
    metadata: row.metadata || {},
    artifactPath: row.artifact_path,
    createdAt: row.created_at,
    updatedAt: row.updated_at,
  }
}

function mapModelVersion(row: any): ModelVersion {
  return {
    id: row.id,
    modelId: row.model_id,
    version: row.version,
    artifactPath: row.artifact_path,
    metadata: row.metadata || {},
    createdAt: row.created_at,
  }
}

