import { Injectable } from "@nestjs/common";
import { ConfigService } from "@nestjs/config";
import { OpenAIWrappedService } from "../../openai-wrapped/openai-wrapped.service";
import { PullRequestGithubEventsVectorService } from "../../timescale/pull_request_github-events_vector.service";
import {
  PullRequestAgentParams,
  SearchAllPrsInDatasetParams,
  SearchAllPrsParams,
  SearchPrsByAuthorInDatasetParams,
  SearchPrsByAuthorParams,
  SearchPrsByRepoNameAndAuthorParams,
  SearchPrsByRepoNameParams,
} from "../schemas/pull-requests.schema";

@Injectable()
export class PullRequestAgent {
  agentSystemMessage: string;

  constructor(
    private configService: ConfigService,
    private openAIWrappedService: OpenAIWrappedService,
    private pullRequestGithubEventsVectorService: PullRequestGithubEventsVectorService
  ) {
    this.agentSystemMessage = this.configService.get("starsearch.pullRequestAgentSystemMessage")!;
  }

  private async searchAllPrs({ question }: SearchAllPrsParams) {
    const queryEmbedding = await this.openAIWrappedService.generateEmbedding(question);

    return this.pullRequestGithubEventsVectorService.cosineSimilarity({
      embedding: queryEmbedding,
      range: 30,
      prevDaysStartDate: 0,
    });
  }

  public searchAllPrsTool = this.openAIWrappedService.makeRunnableToolFunction({
    function: async (params: SearchAllPrsParams) => this.searchAllPrs(params),
    schema: SearchAllPrsParams,
    name: "searchAllPrs",
    description:
      "Searches ALL GitHub pull requests and their context using cosine similarity search based on the input search query. Returns relevant summaries and metadata of pull requests based on the closest neighbor search algorithm. WARNING: this may return irrelevant results if similarity search has no close neighbors. For better and more relevant results, expand search queries at will with related terms. For example, a search query of 'AI' could be expanded to include 'machine learning optimization', 'neural network', and 'model'.",
  });

  private async searchAllPrsInDataset({ question, dataset }: SearchAllPrsInDatasetParams) {
    const queryEmbedding = await this.openAIWrappedService.generateEmbedding(question);

    return this.pullRequestGithubEventsVectorService.cosineSimilarity({
      embedding: queryEmbedding,
      range: 30,
      prevDaysStartDate: 0,
      repoNames: dataset.split(","),
    });
  }

  public searchAllPrsInDatasetTool = this.openAIWrappedService.makeRunnableToolFunction({
    function: async (params: SearchAllPrsInDatasetParams) => this.searchAllPrsInDataset(params),
    schema: SearchAllPrsInDatasetParams,
    name: "searchAllPrsInDataset",
    description:
      "Searches all GitHub pull requests and their context within a given dataset of repositories using cosine similarity search based on the input search query. Returns relevant summaries and metadata of pull requests based on a closest neighbor search algorithm. WARNING: this may return irrelevant results if similarity search has no close neighbors. For better and more relevant results, expand search queries at will with related terms. For example, a search query of 'AI' could be expanded to include 'machine learning optimization', 'neural network', and 'model'.",
  });

  private async searchPrsByRepoName({ question, repoName, range }: SearchPrsByRepoNameParams) {
    const queryEmbedding = await this.openAIWrappedService.generateEmbedding(question);

    return this.pullRequestGithubEventsVectorService.cosineSimilarity({
      embedding: queryEmbedding,
      range,
      prevDaysStartDate: 0,
      repoNames: [repoName],
    });
  }

  public searchPrsByRepoNameTool = this.openAIWrappedService.makeRunnableToolFunction({
    function: async (params: SearchPrsByRepoNameParams) => this.searchPrsByRepoName(params),
    schema: SearchPrsByRepoNameParams,
    name: "searchPrsByRepoName",
    description:
      "Searches GitHub pull requests and their context within a specific repository using cosine similariity search based on the input search query. Returns relevant summaries and metadata of pull requests based on a closest neighbor search algorithm. WARNING: this may return irrelevant results if similairyt search has no close neighbors. For better and more relevant results, expand search queries at will with related terms. For example, a search query of 'AI' could be expanded to include 'machine learning optimization', 'neural network', and 'model'.",
  });

  private async searchPrsByAuthor({ question, author }: SearchPrsByAuthorParams) {
    const queryEmbedding = await this.openAIWrappedService.generateEmbedding(question);

    return this.pullRequestGithubEventsVectorService.cosineSimilarity({
      embedding: queryEmbedding,
      range: 30,
      prevDaysStartDate: 0,
      author,
    });
  }

  public searchPrsByAuthorTool = this.openAIWrappedService.makeRunnableToolFunction({
    function: async (params: SearchPrsByAuthorParams) => this.searchPrsByAuthor(params),
    schema: SearchPrsByAuthorParams,
    name: "searchPrsByAuthor",
    description:
      "Searches GitHub pull requests and their context by a specific PR author using cosine similarity search based on the input search query. Returns relevant summaries and metadata of pull requests based on a closest neighbor search algorithm. WARNING: this may return irrelevant results if similairyt search has no close neighbors. For better and more relevant results, expand search queries at will with related terms. For example, a search query of 'AI' could be expanded to include 'machine learning optimization', 'neural network', and 'model'.",
  });

  private async searchPrsByRepoNameAndAuthor({ question, repoName, author }: SearchPrsByRepoNameAndAuthorParams) {
    const queryEmbedding = await this.openAIWrappedService.generateEmbedding(question);

    return this.pullRequestGithubEventsVectorService.cosineSimilarity({
      embedding: queryEmbedding,
      range: 30,
      prevDaysStartDate: 0,
      repoNames: [repoName],
      author,
    });
  }

  public searchPrsByRepoNameAndAuthorTool = this.openAIWrappedService.makeRunnableToolFunction({
    function: async (params: SearchPrsByRepoNameAndAuthorParams) => this.searchPrsByRepoNameAndAuthor(params),
    schema: SearchPrsByRepoNameAndAuthorParams,
    name: "searchPrsByRepoNameAndAuthor",
    description:
      "Searches GitHub pull requests and their context within a specific repository and by a specific PR author using cosine similairty search based on the input search query. Returns relevant summaries and metadata of pull requests based on a closest neighbor search algorithm. WARNING: this may return irrelevant results if similairyt search has no close neighbors. For better and more relevant results, expand search queries at will with related terms. For example, a search query of 'AI' could be expanded to include 'machine learning optimization', 'neural network', and 'model'.",
  });

  private async searchPrsByAuthorInDataset({ question, author, dataset }: SearchPrsByAuthorInDatasetParams) {
    const queryEmbedding = await this.openAIWrappedService.generateEmbedding(question);

    return this.pullRequestGithubEventsVectorService.cosineSimilarity({
      embedding: queryEmbedding,
      range: 30,
      prevDaysStartDate: 0,
      repoNames: dataset.split(","),
      author,
    });
  }

  public searchPrsByAuthorInDatasetTool = this.openAIWrappedService.makeRunnableToolFunction({
    function: async (params: SearchPrsByAuthorInDatasetParams) => this.searchPrsByAuthorInDataset(params),
    schema: SearchPrsByAuthorInDatasetParams,
    name: "searchPrsByAuthorIndataset",
    description:
      "Searches GitHub pull requests and their context by a specific PR author within a given dataset or repositories using cosine similarity search based on the input search query. Returns relevant summaries and metadata of pull requests based on a closest neighbor search algorithm. WARNING: this may return irrelevant results if similairyt search has no close neighbors. For better and more relevant results, expand search queries at will with related terms. For example, a search query of 'AI' could be expanded to include 'machine learning optimization', 'neural network', and 'model'.",
  });

  async runTools(agentParams: PullRequestAgentParams): Promise<string | null | unknown> {
    const tools = [
      this.searchAllPrsTool,
      this.searchAllPrsInDatasetTool,
      this.searchPrsByRepoNameTool,
      this.searchPrsByAuthorTool,
      this.searchPrsByRepoNameAndAuthorTool,
      this.searchPrsByAuthorInDatasetTool,
    ];

    const runner = this.openAIWrappedService
      .runTools(this.agentSystemMessage, agentParams.prompt, tools)
      .on("message", (msg) => console.log("pull req agent msg", msg))
      .on("functionCall", (functionCall) => console.log("pull req agent functionCall", functionCall))
      .on("functionCallResult", (functionCallResult) =>
        console.log("pull req agent functionCallResult", functionCallResult)
      );

    return runner.finalContent();
  }
}