import z from 'zod';
import {
  authedServiceRequest,
  serviceRequest,
  serviceResponse,
} from '../interfaces/common';
import { metadata } from '../interfaces/Metadata';
import { geolocation } from '../interfaces/Geolocation';
import { zGPTModels } from '../enums/GPTModels';
import { zInstallationChannels } from '../enums/InstallationChannels';
import { zSearchTypes } from '../enums/SearchTypes';
import { Snippet } from './SnippetService';

export const queryEdit = z.object({
  _id: z.string().optional(),
  value: z.string(),
  meta: metadata,
});

export const queryOrigin = z.object({
  channel: zInstallationChannels,
  originId: z.string().optional(),
  originName: z.string().optional(),
  originChannelId: z.string().optional(),
  originChannelName: z.string().optional(),
});

export const query = z.object({
  _id: z.string(),
  distinctId: z.string(),
  organizationId: z.string(),
  botId: z.string(),
  sourceIds: z.array(z.string()),
  snippetIds: z.array(z.string()),
  input: z.string(),
  output: z.string(),
  helpful: z.boolean().nullable(), // user marked as helpful
  mark: z.boolean().nullable(), // admin marked as correct or incorrect
  edits: z.array(queryEdit),
  model: zGPTModels.nullable(),
  embedding: z.array(z.number()).nullable(),
  type: zSearchTypes,
  origin: queryOrigin.nullable(),
  cachedQueryId: z.string().nullable(),
  geolocation: geolocation.nullable(),
  lang: z.string().optional(),
  meta: metadata,
});

export type Query = z.infer<typeof query>;
export type QueryEdit = z.infer<typeof queryEdit>;
export type QueryOrigin = z.infer<typeof queryOrigin>;

export type QueryResolved = Query & {
  snippets: Snippet[];
};
export interface QueryService {
  create(request: CreateQueryRequest): Promise<CreateQueryResponse>;
  setFeedback(request: SetQueryFeedbackRequest): Promise<SetQueryFeedbackResponse>;
  addEdit(request: AddQueryEditRequest): Promise<AddQueryEditResponse>;
  get(request: GetQueryRequest): Promise<GetQueryResponse>;
  list(request: ListQueriesRequest): Promise<ListQueriesResponse>;
  count(request: CountQueriesRequest): Promise<CountQueriesResponse>;
}

/** ******************************************************************************
*  Create Query
******************************************************************************* */

export const createQueryParams = z.object({
  distinctId: z.string(),
  organizationId: z.string(),
  botId: z.string(),
  sourceIds: z.array(z.string()),
  snippetIds: z.array(z.string()),
  input: z.string(),
  output: z.string(),
  model: zGPTModels.nullable(),
  embedding: z.array(z.number()).nullable(),
  type: zSearchTypes,
  origin: queryOrigin.nullable(),
  cachedQueryId: z.string().nullable(),
  geolocation: geolocation.nullable(),
  lang: z.string().optional(),
});

export const createQueryRequest = serviceRequest.merge(
  z.object({
    params: createQueryParams,
  }),
);

export const createQueryResponse = serviceResponse.merge(
  z.object({
    query: query.optional(),
  }),
);

export type CreateQueryParams = z.infer<typeof createQueryParams>;
export type CreateQueryRequest = z.infer<typeof createQueryRequest>;
export type CreateQueryResponse = z.infer<typeof createQueryResponse>;

/** ******************************************************************************
*  Set Query Feedback
******************************************************************************* */

export const setQueryFeedbackParams = z.object({
  queryId: z.string(),
  helpful: z.boolean().nullable(),
});

export const setQueryFeedbackRequest = serviceRequest.merge(
  z.object({
    params: setQueryFeedbackParams,
  }),
);

export const setQueryFeedbackResponse = serviceResponse.merge(
  z.object({
    query: query.optional(),
  }),
);

export type SetQueryFeedbackParams = z.infer<typeof setQueryFeedbackParams>;
export type SetQueryFeedbackRequest = z.infer<typeof setQueryFeedbackRequest>;
export type SetQueryFeedbackResponse = z.infer<typeof setQueryFeedbackResponse>;

/** ******************************************************************************
*  Set Query Correction
******************************************************************************* */

export const addQueryEditParams = z.object({
  _id: z.string(),
  mark: z.boolean().nullable(),
  edit: z.string().nullable(),
});

export const addQueryEditRequest = authedServiceRequest.merge(
  z.object({
    params: addQueryEditParams,
  }),
);

export const addQueryEditResponse = serviceResponse.merge(
  z.object({
    query: query.optional(),
  }),
);

export type AddQueryEditParams = z.infer<typeof addQueryEditParams>;
export type AddQueryEditRequest = z.infer<typeof addQueryEditRequest>;
export type AddQueryEditResponse = z.infer<typeof addQueryEditResponse>;

/** ******************************************************************************
*  Get Query
******************************************************************************* */

export const getQueryParams = z.object({
  queryId: z.string(),
});

export const getQueryRequest = serviceRequest.merge(
  z.object({
    params: getQueryParams,
  }),
);

export const getQueryResponse = serviceResponse.merge(
  z.object({
    query: query.nullable().optional(),
  }),
);

export type GetQueryParams = z.infer<typeof getQueryParams>;
export type GetQueryRequest = z.infer<typeof getQueryRequest>;
export type GetQueryResponse = z.infer<typeof getQueryResponse>;

/** ******************************************************************************
*  List Queries
******************************************************************************* */

export const listQueriesParams = z.object({
  botId: z.string().optional(),
  sourceIds: z.array(z.string()).optional(),
  search: z.string().optional(),
  offset: z.number().optional(),
  pageSize: z.number().optional(),
  isSuper: z.boolean().optional(),
});

export const listQueriesRequest = authedServiceRequest.merge(
  z.object({
    params: listQueriesParams,
  }),
);

export const listQueriesResponse = serviceResponse.merge(
  z.object({
    queries: z.array(query).optional(),
    total: z.number().optional(),
  }),
);

export type ListQueriesParams = z.infer<typeof listQueriesParams>;
export type ListQueriesRequest = z.infer<typeof listQueriesRequest>;
export type ListQueriesResponse = z.infer<typeof listQueriesResponse>;

/** ******************************************************************************
*  Count Queries
******************************************************************************* */

export const countQueriesParams = z.object({
  botId: z.string().optional(),
  sourceIds: z.array(z.string()).optional(),
  search: z.string().optional(),
  isSuper: z.boolean().optional(),
});

export const countQueriesRequest = authedServiceRequest.merge(
  z.object({
    params: countQueriesParams,
  }),
);

export const countQueriesResponse = serviceResponse.merge(
  z.object({
    count: z.number().optional(),
  }),
);

export type CountQueriesParams = z.infer<typeof countQueriesParams>;
export type CountQueriesRequest = z.infer<typeof countQueriesRequest>;
export type CountQueriesResponse = z.infer<typeof countQueriesResponse>;