const { LangChainTracer } = require("@langchain/core/tracers/tracer_langchain");
const axios = require("axios");

/**
 * @typedef {Object} JLINCTracerFields
 * @property {string} [dataStoreApiUrl]
 * @property {string} [dataStoreApiKey]
 * @property {string} [archiveApiUrl]
 * @property {string} [archiveApiKey]
 * @property {string} [systemPrefix]
 * @property {string|null} [agreementId]
 * @property {boolean} [debug]
 */

/**
 * @typedef {Object} SerializedRun
 * @property {string} [name]
 * @property {number} [start_time]
 * @property {number} [end_time]
 * @property {string} [trace_id]
 * @property {any} [inputs]
 * @property {any} [outputs]
 * @property {{ metadata?: { ls_provider?: string, ls_model_name?: string } }} [extra]
 */

/**
 * @extends {LangChainTracer}
 */
class JLINCTracer extends LangChainTracer {
  /**
   * @param {JLINCTracerFields} [fields={}]
   */
  constructor(fields = {}) {
    super(fields);
    const defaultDataStoreApiUrl = "https://api-test.jlinc.io";
    const defaultDataStoreApiKey = "";
    const defaultArchiveApiUrl = "https://archive-test.jlinc.io";
    const defaultArchiveApiKey = "";
    const defaultSystemPrefix = "MyProject";

    /** @type {string} */
    this.dataStoreApiUrl = fields.dataStoreApiUrl ?? defaultDataStoreApiUrl;
    /** @type {string} */
    this.dataStoreApiKey = fields.dataStoreApiKey ?? defaultDataStoreApiKey;
    /** @type {string} */
    this.archiveApiUrl = fields.archiveApiUrl ?? defaultArchiveApiUrl;
    /** @type {string} */
    this.archiveApiKey = fields.archiveApiKey ?? defaultArchiveApiKey;
    /** @type {string} */
    this.systemPrefix = this._sanitize(fields.systemPrefix ?? defaultSystemPrefix);
    /** @type {string|null} */
    this.agreementId = fields.agreementId ?? null;
    /** @type {boolean} */
    this.debug = fields.debug ?? false;

    /** @type {string|null} */
    this.domain = null;
    /** @type {Set<string>} */
    this.entities = new Set();
  }

  /**
   * @param {any} run
   * @returns {Promise<void>}
   */
  async onRunCreate(run) { }

  /**
   * @param {any} run
   * @returns {Promise<void>}
   */
  async onRunUpdate(run) { }

  /**
   * @returns {Promise<void>}
   */
  async getDefaultProjectName() {
    return this.systemPrefix;
  }

  /**
   * @param {SerializedRun} serialized
   * @returns {Promise<void>}
   */
  async onLLMStart(serialized) {
    if (this.debug) {
      console.log(`\nJLINC: onLLMStart`);
      console.log(`---------------------------------------------`);
    }
    await this._jlincTrace({
      event: "llm_start",
      name: serialized?.name || "unknown",
      ts: serialized?.start_time,
      provider: serialized?.extra?.metadata,
      traceId: serialized?.trace_id,
      inputs: serialized?.inputs,
    }, 'out');
  }

  /**
   * @param {SerializedRun} serialized
   * @returns {Promise<void>}
   */
  async onLLMEnd(serialized) {
    if (this.debug) {
      console.log(`\nJLINC: onLLMEnd`);
      console.log(`---------------------------------------------`);
    }
    await this._jlincTrace({
      event: "llm_end",
      name: serialized?.name || "unknown",
      ts: serialized?.end_time,
      provider: serialized?.extra?.metadata,
      traceId: serialized?.trace_id,
      outputs: serialized?.outputs,
    }, 'in');
  }

  /**
   * @param {SerializedRun} serialized
   * @returns {Promise<void>}
   */
  async onToolStart(serialized) {
    if (this.debug) {
      console.log(`\nJLINC: onToolStart`);
      console.log(`---------------------------------------------`);
    }
    await this._jlincTrace({
      event: "tool_start",
      name: serialized?.name || "unknown",
      ts: serialized?.start_time,
      provider: serialized?.extra?.metadata,
      traceId: serialized?.trace_id,
      inputs: serialized?.inputs,
    }, 'out');
  }

  /**
   * @param {SerializedRun} serialized
   * @returns {Promise<void>}
   */
  async onToolEnd(serialized) {
    if (this.debug) {
      console.log(`\nJLINC: onToolEnd`);
      console.log(`---------------------------------------------`);
    }
    await this._jlincTrace({
      event: "tool_end",
      name: serialized?.name || "unknown",
      ts: serialized?.end_time,
      provider: serialized?.extra?.metadata,
      traceId: serialized?.trace_id,
      outputs: serialized?.outputs,
    }, 'in');
  }

  /**
   * @param {string} dir
   * @param {any} payload
   * @returns {Promise<any>}
   */
  async _postToApi(dir, payload) {
    const response = await axios.post(
      `${this.dataStoreApiUrl}/api/v1/data/${dir}`,
      payload,
      {
        headers: {
          'Authorization': `Bearer ${this.dataStoreApiKey}`,
        }
      }
    );
    return response.data;
  }

  /**
   * @param {string} entityShortName
   * @returns {Promise<void>}
   */
  async _setEntity(entityShortName) {
    if (!this.entities.has(entityShortName)) {
      let recipient = null;
      try {
        recipient = await this._postToApi(`entity/get`, { shortName: entityShortName });
      } catch (e) {
        recipient = await this._postToApi(`entity/create`, { shortName: entityShortName });
      }
      this.entities.add(entityShortName);
    }
  }

  /**
   * @param {string} input
   * @returns {string}
   */
  _sanitize(input) {
    return input.replace(/[^a-zA-Z0-9._-]/g, '_');
  }

  /**
   * @param {any} payload
   * @param {'in' | 'out'} direction
   * @returns {Promise<void>}
   */
  async _jlincTrace(payload, direction) {
    try {
      if (!this.domain) {
        const domains = await this._postToApi(`entity/domains/get`, {});
        this.domain = domains[0];
      }
      let entityShortName = `${this.systemPrefix}-${this._sanitize(payload.name)}`;
      payload?.provider?.ls_provider && (entityShortName += `-${this._sanitize(payload.provider.ls_provider)}`);
      payload?.provider?.ls_model_name && (entityShortName += `-${this._sanitize(payload.provider.ls_model_name)}`);
      entityShortName += `@${this.domain}`;
      await this._setEntity(entityShortName);
      const systemShortName = `${this.systemPrefix}-system@${this.domain}`;
      await this._setEntity(systemShortName);
      const recipientShortName = direction === 'in' ? systemShortName : entityShortName;
      const senderShortName = direction === 'out' ? systemShortName : entityShortName;
      const data = {
        type: 'data',
        senderShortName,
        recipientShortName,
        agreementId: this.agreementId,
        data: payload,
        archive: {
          url: this.archiveApiUrl,
          key: this.archiveApiKey,
        }
      };
      const event = await this._postToApi(`event/produce`, data);
      if (this.debug) {
        console.log(JSON.stringify({ event }, null, 2));
      }
    } catch (error) {
      console.error("[Tracer] Failed to log event:", error.message);
    }
  }
 
}

module.exports = { JLINCTracer }