import { GraphQLClient } from 'graphql-request';
import {
  GetUserQueryVariables,
  CreateUserMutationVariables,
  UpdateUserMutationVariables,
  DeleteUserMutationVariables,
  PredictionWhereUniqueInput,
  PredictionWhereInput,
  MetadataInput
} from '../gql/graphql';
import { graphql } from '../gql';
import { BaseFunctions } from './BaseFunctions';
import { ListPredictionsQueryVariables,  } from '../gql/graphql';
import { OpenAiKwargs } from '../types';



class PredictionFunctions extends BaseFunctions {

  constructor(client: GraphQLClient) {
    super(client);
  }

  // a method to create a prediction
  async create(openaiKwargs: OpenAiKwargs, prompt: string, modelBackendName: string, useCache: boolean = true, labelParser?: string, metadata?: MetadataInput[]) {
    const query = graphql(`
      query Prediction($openaiKwargs: OpenAIKwargs!, $prompt: String!, $modelBackendName: String!, $useCache: Boolean!, $labelParser: String, $metadata: [MetadataInput!]) {
        prediction(openaiKwargs: $openaiKwargs, prompt: $prompt, modelBackendName: $modelBackendName, useCache: $useCache, labelParser: $labelParser, metadata: $metadata) {
          completion
          id
          labels {
            id
            confidence
            name
            value
          }
          prompt
          labelParser
          metadata {
            key
            value
          }
        }
      }
    `);
    const vars = { openaiKwargs, prompt, modelBackendName, useCache, labelParser, metadata };
    return this.request(query, vars);
  }

  async countPredictions(where?: PredictionWhereUniqueInput) {
    const query = graphql(`
      query CountPredictions($where: PredictionWhereInput) {
        aggregatePrediction(where: $where) {
          _count {
            _all
          }
        }
      }
    `);
    let vars: any = {}
    if (where) {
      vars.where = where
    }
    return this.request(query, vars);
  }

  async getLabelsForPrediction(predictionId: number) {
    const query = graphql(`
      query GetLabelsForPrediction($where: LabelWhereInput) {
        labels(where: $where) {
          id
          value
          confidence
          predictionId
        }
      }
    `);
    return this.request(query, { where: { predictionId: { equals: predictionId } } });
  }

  async listPredictions(options: ListPredictionsQueryVariables) {
    const query = graphql(`
    query ListPredictions($where: PredictionWhereInput, $orderBy: [PredictionOrderByWithRelationInput!], $cursor: PredictionWhereUniqueInput, $take: Int, $skip: Int, $distinct: [PredictionScalarFieldEnum!]) {
      predictions(where: $where, orderBy: $orderBy, cursor: $cursor, take: $take, skip: $skip, distinct: $distinct) {
        id
        completion
        createdAt
        _count {
          labels
          datasets
        }
        labels {
          id
          name
          value
          confidence
        }
        prompt
        requestArgs
        response
        updatedAt
        userId
        modelId
      }
    }
    `);
    return this.request(query, options);
  }

  async getPrediction(predictionId: number) {
    const query = graphql(`
        query GetPrediction($id: Int!) {
          getPrediction(where: { id: $id }) {
            completion
            createdAt
            id
            _count {
              labels
              datasets
            }
            model {
              name
            }
            prompt
          }
        }
      `);

    return this.request(query, { id: predictionId });
  }

  async addPredictionToDataset(predictionId: number, datasetId: number) {
    const mutation = graphql(`
      mutation UpdateOnePrediction($predictionId: Int!, $datasetId: Int!) {
        updateOnePrediction(where: { id: $predictionId }, data: { datasets: { connect: { id: $datasetId } } }) {
          id
          _count {
            datasets
          }
        }
      }
    `)

    return this.request(mutation, { predictionId, datasetId })
  }

  async addPredictionsToDataset(predictionIds: number[], datasetId: number) {
    const mutation = graphql(`
      mutation Mutation($datasetId: Int!, $predictionIds: [PredictionWhereUniqueInput!]!) {
        updateOneDataset(where: { id: $datasetId }, data: { predictions: { connect: $predictionIds } }) {
          id
          _count {
            predictions
          }
        }
      }
    `)

    return this.request(mutation, { predictionIds: predictionIds.map(id => ({ id })), datasetId })
  }
}


export { PredictionFunctions };
