import { inject, Injectable } from '@angular/core';
import { renderMarkdown, SDKService } from '@flaps/core';
import { ModalConfig, ModalRef, ModalService } from '@guillotinaweb/pastanaga-angular';
import { BehaviorSubject, forkJoin, Observable, of, switchMap, take, tap } from 'rxjs';
import {
  Ask,
  ChatOptions,
  GenerativeProvider,
  IErrorResponse,
  ModelInfo,
  NUCLIA_STANDARD_SEARCH_CONFIG,
  Prompts,
  TokenConsumption,
  Widget,
} from '@nuclia/core';
import { LoadingDialogComponent } from './loading-dialog';
import { catchError, filter, map } from 'rxjs/operators';
import { SearchWidgetService } from '../search-widget';
import { GENERATIVE_MODEL_KEY, RequestConfig, RequestConfigAndQueries, ResultEntry } from './rag-lab.models';
import { TranslateService } from '@ngx-translate/core';

@Injectable({
  providedIn: 'root',
})
export class RagLabService {
  private sdk = inject(SDKService);
  private modalService = inject(ModalService);
  private translate = inject(TranslateService);
  private searchWidgetService = inject(SearchWidgetService);

  private _loadingModal?: ModalRef;
  private _progress$ = new BehaviorSubject<number | null>(null);
  private _promptLabResults = new BehaviorSubject<
    {
      query: string;
      prompt?: Prompts;
      results: ResultEntry[];
    }[]
  >([]);
  private _ragLabResults = new BehaviorSubject<
    {
      query: string;
      prompt?: Prompts;
      results: ResultEntry[];
    }[]
  >([]);

  private _kbConfigBackup = new BehaviorSubject<{ [id: string]: any } | null>(null);
  private _generativeProviders = new BehaviorSubject<GenerativeProvider[] | null>(null);
  private _generativeModelMap = new BehaviorSubject<{ [value: string]: string } | null>(null);
  private _searchConfigurations = new BehaviorSubject<Widget.TypedSearchConfiguration[]>([]);

  kbConfigBackup = this._kbConfigBackup.asObservable();
  generativeProviders = this._generativeProviders.asObservable();
  generativeModelMap = this._generativeModelMap.asObservable();
  searchConfigurations = this._searchConfigurations.asObservable();
  promptLabResults = this._promptLabResults.asObservable();
  ragLabResults = this._ragLabResults.asObservable();

  private getModelName(model: string) {
    return this._generativeModelMap.value?.[model] || model;
  }

  loadKbConfigAndModels(): Observable<void> {
    return this.sdk.currentKb.pipe(
      take(1),
      switchMap((kb) =>
        forkJoin([
          kb.getConfiguration(),
          kb.getLearningSchema(),
          kb.getGenerativeProviders(),
          this.searchWidgetService.supportedSearchConfigurations.pipe(take(1)),
        ]),
      ),
      map(([config, schema, providers, savedConfigurations]) => {
        this._searchConfigurations.next([{ ...NUCLIA_STANDARD_SEARCH_CONFIG }].concat(savedConfigurations));
        this._kbConfigBackup.next(config);
        // Only models present in both schema and providers are allowed
        const allowedModels = (schema[GENERATIVE_MODEL_KEY]?.options || []).map((option) => option.value);
        const allowedProviders = Object.entries(providers)
          .map(([providerKey, provider]) => ({
            ...provider,
            models: Object.fromEntries(
              Object.entries(provider.models)
                .map(([key, value]) =>
                  // The key and name of custom models are swapped to keep consistency with the other models
                  providerKey === 'default' ? [value.name, { ...value, name: key }] : [key, value],
                )
                .map((value) => value as [string, ModelInfo])
                .filter(([key]) => allowedModels.includes(key) && key !== 'generative-multilingual-2023'),
            ),
          }))
          .filter((provider) => Object.keys(provider.models || {}).length > 0);
        this._generativeProviders.next(allowedProviders);
        this._generativeModelMap.next(
          allowedProviders.reduce(
            (acc, provider) => {
              Object.entries(provider.models).forEach(([key, model]) => {
                acc[key] = provider.title + ' - ' + model.title;
              });
              return acc;
            },
            {} as { [value: string]: string },
          ),
        );
      }),
    );
  }

  generate(requestConfigs: RequestConfigAndQueries[], forTab: 'prompt' | 'rag') {
    this._progress$.next(null);
    const requestCount = requestConfigs.reduce((count, requestConfig) => {
      return count + requestConfig.queries.length;
    }, 0);

    if (forTab === 'prompt') {
      this._promptLabResults.next([]);
    } else if (forTab === 'rag') {
      this._ragLabResults.next([]);
    }

    this._loadingModal = this.modalService.openModal(
      LoadingDialogComponent,
      new ModalConfig<{ title: string; description: string; progress: Observable<number | null> }>({
        dismissable: false,
        data: {
          title: 'rag-lab.common.generate.loading.title',
          description: 'rag-lab.common.generate.loading.description',
          progress: this._progress$
            .asObservable()
            .pipe(map((progress) => Math.ceil(((progress || 0) * 100) / requestCount))),
        },
      }),
    );
    this._progress$.pipe(filter((progress) => progress === requestCount)).subscribe(() => this._loadingModal?.close());

    return this._generateResults(requestConfigs, forTab);
  }

  private _generateResults(requestConfigs: RequestConfigAndQueries[], forTab: 'prompt' | 'rag') {
    const requests: { query: string; options: ChatOptions }[] = requestConfigs.reduce(
      (requests, config) => {
        const { queries, ...options } = config;
        return requests.concat(queries.map((query) => ({ query, options })));
      },
      [] as { query: string; options: ChatOptions }[],
    );

    return this.sdk.currentKb.pipe(
      take(1),
      switchMap((kb) =>
        forkJoin(
          requests.map(({ query, options }) =>
            kb
              .ask(query, undefined, undefined, {
                ...options,
                synchronous: true,
                debug: true,
                show_consumption: true,
              })
              .pipe(
                switchMap((answer) => {
                  if (answer.type === 'error') {
                    this._addResult(
                      {
                        query,
                        options,
                        result:
                          answer.status === -1
                            ? this.translate.instant('rag-lab.common.results.error-llm')
                            : this.translate.instant('rag-lab.common.results.error'),
                      },
                      forTab,
                    );
                    return of(null);
                  } else {
                    return renderMarkdown(answer.text).pipe(
                      tap((rendered) =>
                        this._addResult(
                          {
                            query,
                            result: answer,
                            options,
                            rendered,
                            tokens:
                              answer.consumption && (answer.consumption?.normalized_tokens.input || 0) > 0
                                ? answer.consumption.normalized_tokens
                                : answer.consumption?.customer_key_tokens,
                            timings: answer.metadata?.timings,
                          },
                          forTab,
                        ),
                      ),
                    );
                  }
                }),
                catchError(() => {
                  this._addResult(
                    {
                      query,
                      options,
                      result: this.translate.instant('rag-lab.common.results.error'),
                    },
                    forTab,
                  );
                  return of(null);
                }),
              ),
          ),
        ),
      ),
    );
  }

  private _addResult(
    entry: {
      query: string;
      result: Ask.Answer | IErrorResponse | string;
      options: RequestConfig;
      rendered?: string;
      tokens?: TokenConsumption;
      timings?: Ask.AskTimings;
    },
    forTab: 'prompt' | 'rag',
  ) {
    const { query, options, result, tokens, timings } = entry;
    const configId = options.searchConfigId;
    const model = options.generative_model as string;
    const modelName = this.getModelName(model);
    const results = forTab === 'prompt' ? this._promptLabResults : this._ragLabResults;
    const entries = results.value;
    const entryIndex = entries.findIndex((entry) => entry.query === query);
    const queryEntry = entryIndex > -1 ? entries[entryIndex] : null;
    const answer =
      typeof result === 'string'
        ? result
        : result.type === 'answer'
          ? (result as Ask.Answer).text
          : `Error: ${result.detail}`;
    if (queryEntry) {
      queryEntry.results.push({ configId, model, modelName, answer, rendered: entry.rendered, tokens, timings });
      entries.splice(entryIndex, 1, queryEntry);
      results.next([...entries]);
    } else {
      results.next(
        entries.concat([
          {
            query,
            prompt: options.prompt as Prompts,
            results: [{ configId, model, modelName, answer, rendered: entry.rendered, tokens, timings }],
          },
        ]),
      );
    }

    this._progress$.next((this._progress$.value || 0) + 1);
  }

  downloadPromptLabCsv(generativeModels: string[]) {
    if (this._promptLabResults.value.length === 0) {
      return;
    }
    const modelsNames = generativeModels.map((model) => this.getModelName(model));
    let csv = `Query,${modelsNames.join(',')},Metadata`;
    this._promptLabResults.value.forEach(({ query, prompt, results }) => {
      const answers = generativeModels
        .map((model) => {
          const modelResult = results.find((result) => result.model === model);
          return `"${this.formatCellValue(modelResult?.answer || 'No answer')}"`;
        })
        .join(',');
      const metadata = `"${this.formatCellValue(this.getMetadata(modelsNames, results))}"`;
      csv += `\n"Query: ${this.formatCellValue(query)}${
        (prompt?.user ? '\nPrompt: ' + this.formatCellValue(prompt.user) : '') +
        (prompt?.system ? '\nSystem prompt: ' + this.formatCellValue(prompt.system) : '')
      }",${answers},${metadata}`;
    });
    const filename = `${new Date().toISOString().split('T')[0]}_Nuclia_models_comparison.csv`;
    this.generateCsv(csv, filename);
  }

  downloadRagLabCsv(configs: { configId: string; generativeModel: string }[]) {
    if (this._ragLabResults.value.length === 0) {
      return;
    }
    const configIds = configs.map((config) => `${config.configId} (${config.generativeModel})`);
    let csv = `Query,${configIds.join(',')},Metadata`;
    this._ragLabResults.value.forEach(({ query, prompt, results }) => {
      const answers = configs
        .map((config) => {
          const modelResult = results.find((result) => result.configId === config.configId);
          return `"${this.formatCellValue(modelResult?.answer || 'No answer')}"`;
        })
        .join(',');
      const metadata = `"${this.formatCellValue(this.getMetadata(configIds, results))}"`;
      csv += `\n"Query: ${this.formatCellValue(query)}${
        (prompt?.user ? '\nPrompt: ' + this.formatCellValue(prompt.user) : '') +
        (prompt?.system ? '\nSystem prompt: ' + this.formatCellValue(prompt.system) : '')
      }",${answers},${metadata}`;
    });
    const filename = `${new Date().toISOString().split('T')[0]}_Nuclia_RAG_comparison.csv`;
    this.generateCsv(csv, filename);
  }

  private generateCsv(csv: string, filename: string) {
    const blob = new Blob([csv], { type: 'text/csv;charset=utf-8' });
    const link = document.createElement('a');
    if (link.download !== undefined) {
      const url = URL.createObjectURL(blob);
      link.setAttribute('href', url);
      link.setAttribute('download', filename);
      link.click();
      URL.revokeObjectURL(url);
    }
  }

  private getMetadata(columns: string[], results: ResultEntry[]) {
    return results
      .map((result, i) => {
        const tokens = result?.tokens
          ? [
              `${this.translate.instant('rag-lab.common.results.input-tokens')}: ${result.tokens.input}`,
              `${this.translate.instant('rag-lab.common.results.output-tokens')}: ${result.tokens.output}`,
              `${this.translate.instant('rag-lab.common.results.image-tokens')}: ${result.tokens.image}`,
            ]
          : [];
        const timings = result?.timings
          ? [
              `${this.translate.instant('rag-lab.common.results.first-word-time')}: ${
                result.timings?.generative_first_chunk?.toFixed(2) || '-'
              }s`,
              `${this.translate.instant('rag-lab.common.results.total-time')}: ${
                result.timings?.generative_total?.toFixed(2) || '-'
              }s`,
            ]
          : [];
        return `${columns[i]} | ${tokens.concat(timings).join(' | ')}`;
      })
      .join('\n');
  }

  private formatCellValue(value: string): string {
    return value.trim().replace(/"/g, '""');
  }
}
