import { GraphQLClient } from 'graphql-request';
import { graphql } from '../gql';
import { BaseFunctions } from './BaseFunctions';
import { MutationCreateOneDatasetArgs } from '../gql/graphql';
import { ListDatasetsQueryVariables } from '../gql/graphql';

class DatasetFunctions extends BaseFunctions {

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

    async listDatasets(options: ListDatasetsQueryVariables) {
        const query = graphql(`
            query ListDatasets ($where: DatasetWhereInput, $skip: Int, $take: Int, $cursor: DatasetWhereUniqueInput, $distinct: [DatasetScalarFieldEnum!]) {
                datasets(where: $where, skip: $skip, take: $take, orderBy: { createdAt: desc }, cursor: $cursor, distinct: $distinct) {
                    id
                    name
                    createdAt
                    updatedAt
                    _count {
                        predictions
                        trainingPredictions
                    }
                }
            }
        `);

        return this.request(query, options);
    }

    async getDataset(datasetId: number) {
        const query = graphql(`
            query GetDataset($datasetId: Int!) {
                dataset(where: { id: $datasetId }) {
                _count {
                    predictions
                    trainingPredictions
                }
                name
                id
                updatedAt
                createdAt
                }
            }
        `);

        return this.request(query, { datasetId });
    }

    async createDataset(data: MutationCreateOneDatasetArgs['data']) {
        const mutation = graphql(`
            mutation CreateDataset($data: DatasetCreateInput!) {
                createOneDataset(data: $data) {
                id
                name
                }
            }
        `);

        return this.request(mutation, { data });
    }

    async getDatasetCSV(datasetId: number) {
        const query = graphql(`
            query GetDatasetCSV($datasetId: Float!) {
                datasetCSV(datasetId: $datasetId)
            }
        `);
        return this.request(query, { datasetId });
    } 
}

export { DatasetFunctions };
