/*
 * Copyright (C) 2023 Amazon.com, Inc. or its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { z } from 'zod';
import { CloudEvent } from '@project-lakechain/sdk';
import { Prompt } from './prompt';
import { OntologyClassifier } from './classifier';

/**
 * System prompt.
 */
const SYSTEM_PROMPT = `
You are a metadata extractor that receives documents and extract structured metadata
from them in the form of a JSON object.
`;

/**
 * User prompt.
 */
const USER_PROMPT = `
Analyze the provided documents step-by-step and extract structured metadata from them.
Below is the schema of the structured metadata that you must extract.

<json>
  {
    "type": "book | article | paper | report | invoice | etc",
    "createdAt": "YYYY-MM-DDTHH:MM:SSZ",
    "updatedAt": "YYYY-MM-DDTHH:MM:SSZ",
    "title": "The title of the document",
    "description": "A description of the document",
    "keywords": ["Keyword1", "Keyword2", "Keyword3"],
    "topics": ["Topic1", "Topic2", "Topic3"],
    "language": "en | fr | de | es | etc",
    "authors": ["John Doe", "Jane Doe"],
    "publisher": {
      "type": "company | organization | institution | etc",
      "name": "The name of the publisher"
    }
  }
</json>

If you cannot extract a specific field, you can omit it from the JSON object.
A few notes on each fields:
- The createdAt and updatedAt fields must be in ISO 8601 format.
- The language field must be a two-letter language code.
- Keywords are the most prominent contextual words in the document.
- Topics represent the main themes of the document.
`;

/**
 * Default ontology classifier properties.
 */
export const DefaultOntologyClassifierPropsSchema = z.object({

  /**
   * A unique identifier for the classifier.
   */
  classifierType: z.literal('DEFAULT')
});

// The type of the `DefaultOntologyClassifierProps` schema.
export type DefaultOntologyClassifierProps = z.infer<typeof DefaultOntologyClassifierPropsSchema>;

/**
 * The custom ontology classifier builder.
 */
export class DefaultOntologyClassifierBuilder {

  /**
   * The image variation task properties.
   */
  private props: Partial<DefaultOntologyClassifierProps> = {
    classifierType: 'DEFAULT',
  };

  public build(): DefaultOntologyClassifier {
    return (DefaultOntologyClassifier.from(this.props));
  }
}

/**
 * The custom ontology classifier.
 */
export class DefaultOntologyClassifier implements OntologyClassifier {

  /**
   * The `DefaultOntologyClassifier` Builder.
   */
  public static readonly Builder = DefaultOntologyClassifierBuilder;

  /**
   * Creates a new instance of the `DefaultOntologyClassifier` class.
   * @param props the task properties.
   */
  constructor(public props: DefaultOntologyClassifierProps) {}

  /**
   * @returns the type of the classifier.
   */
  public classifierType(): string {
    return (this.props.classifierType);
  }

  /**
   * Creates a new instance of the `DefaultOntologyClassifier` class.
   * @param props the task properties.
   * @returns a new instance of the `DefaultOntologyClassifier` class.
   */
  public static from(props: any) {
    return (new DefaultOntologyClassifier(
      DefaultOntologyClassifierPropsSchema.parse(props)
    ));
  }

  /**
   * @returns the prompts to execute to extract the
   * semantic ontology from documents.
   */
  public getPrompts(): Prompt[] {
    return ([{
      type: 'default',
      userPrompt: USER_PROMPT,
      systemPrompt: SYSTEM_PROMPT
    }]);
  }

  /**
   * Updates the given cloud event based on the answer generated by the LLM
   * and the prompt.
   * @param json the JSON object extracted from the LLM.
   * @param _ the prompt used to generate the JSON object.
   * @param event the cloud event to update.
   * @returns the updated cloud event.
   */
  public update(json: any, _: Prompt, event: CloudEvent): Promise<CloudEvent> {
    const metadata = event.data().metadata();

    // Document type.
    if (json.type) {
      metadata.type = json.type;
    }

    // Dates.
    if (json.createdAt) {
      metadata.createdAt = json.createdAt;
    }
    if (json.updatedAt) {
      metadata.updatedAt = json.updatedAt;
    }

    // Title.
    if (json.title) {
      metadata.title = json.title;
    }

    // Description.
    if (json.description) {
      metadata.description = json.description;
    }

    // Keywords.
    if (Array.isArray(json.keywords)) {
      if (!metadata.keywords?.length) {
        metadata.keywords = json.keywords;  
      } else {
        metadata.keywords.push(...json.keywords);
      }
    }

    // Topics.
    if (Array.isArray(json.topics) && json.topics.length) {
      if (!metadata.topics?.length) {
        metadata.topics = json.topics;  
      } else {
        metadata.topics.push(...json.topics);
      }
    }

    // Authors.
    if (Array.isArray(json.authors) && json.authors.length) {
      metadata.authors = json.authors;
    }

    // Language.
    if (json.language) {
      metadata.language = json.language;
    }

    return (Promise.resolve(event));
  }

  /**
   * @returns the JSON representation of the classifier.
   */
  public toJSON() {
    return (this.props);
  }
}
