// SPDX-License-Identifier: MIT
// Copyright contributors to the openassistant project

import {
  OpenAssistantTool,
  OpenAssistantExecuteFunctionResult,
} from '@openassistant/utils';
import { z } from 'zod';
import { createWeights, WeightsMeta, CreateWeightsProps } from '@geoda/core';
import { WeightsProps, GetGeometries } from '../types';

// global variable to store the created weights
/**
 * @internal
 */
export const globalWeightsData: Record<string, WeightsProps> = {};

export type SpatialWeightsFunctionArgs = z.ZodObject<{
  datasetName: z.ZodString;
  type: z.ZodEnum<['knn', 'queen', 'rook', 'threshold']>;
  k: z.ZodOptional<z.ZodNumber>;
  orderOfContiguity: z.ZodOptional<z.ZodNumber>;
  includeLowerOrder: z.ZodOptional<z.ZodBoolean>;
  precisionThreshold: z.ZodOptional<z.ZodNumber>;
  distanceThreshold: z.ZodOptional<z.ZodNumber>;
  isMile: z.ZodOptional<z.ZodBoolean>;
  useCentroids: z.ZodOptional<z.ZodBoolean>;
  mapBounds: z.ZodOptional<z.ZodArray<z.ZodNumber>>;
}>;

export type SpatialWeightsLlmResult = {
  success: boolean;
  weightsId?: string;
  weightsMeta?: WeightsMeta;
  result?: string;
  error?: string;
  instruction?: string;
};

export type SpatialWeightsAdditionalData = {
  weightsId: string;
  [id: string]:
    | string
    | {
        type: 'weights';
        content: {
          weights: number[][];
          weightsMeta: WeightsMeta;
        };
      };
};

/**
 * ## spatialWeights Tool
 *
 * This tool creates spatial weights matrices for spatial analysis.
 * Spatial weights define the spatial relationships between observations and are essential for spatial autocorrelation analysis.
 *
 * ### Spatial Weights Methods
 *
 * The tool supports various spatial weights methods:
 * - **queen**: Queen contiguity (shares edge or vertex)
 * - **rook**: Rook contiguity (shares edge only)
 * - **knn**: K-nearest neighbors
 * - **threshold**: Distance-based weights
 *
 * ### Parameters
 * - `datasetName`: Name of the dataset containing the geometries
 * - `type`: Spatial weights method to use (see above)
 * - `k`: Number of neighbors for knn method (optional)
 * - `orderOfContiguity`: Order of contiguity for queen/rook (optional)
 * - `includeLowerOrder`: Include lower order contiguity (optional)
 * - `precisionThreshold`: Precision threshold for geometry operations (optional)
 * - `distanceThreshold`: Distance threshold for threshold method (optional)
 * - `isMile`: Use miles instead of kilometers for distance (optional)
 * - `useCentroids`: Use centroids instead of full geometries (optional)
 * - `mapBounds`: Map bounds for visualization (optional)
 *
 * **Example user prompts:**
 * - "Create spatial weights using queen contiguity"
 * - "Generate k-nearest neighbors weights with k=5"
 * - "Make distance-based weights with 10km threshold"
 *
 * ### Example
 * ```typescript
 * import { spatialWeights } from "@openassistant/geoda";
 * import { convertToVercelAiTool } from "@openassistant/utils";
 *
 * const weightsTool = {
 *   ...spatialWeights,
 *   context: {
 *     getGeometries: async (datasetName: string) => {
 *       // Implementation to retrieve geometries from your data source
 *       return geometries;
 *     },
 *   },
 * };
 *
 * const result = await generateText({
 *   model: openai('gpt-4.1', { apiKey: key }),
 *   prompt: 'Create spatial weights using queen contiguity',
 *   tools: { spatialWeights: convertToVercelAiTool(weightsTool) },
 * });
 * ```
 *
 * :::note
 * The spatialWeights tool should always be used with other spatial analysis tools like globalMoran, lisa, etc.
 * :::
 */
export const spatialWeights: OpenAssistantTool<
  SpatialWeightsFunctionArgs,
  SpatialWeightsLlmResult,
  SpatialWeightsAdditionalData,
  SpatialWeightsFunctionContext
> = {
  name: 'spatialWeights',
  description: 'Create a spatial weights.',
  parameters: z.object({
    datasetName: z.string(),
    type: z.enum(['knn', 'queen', 'rook', 'threshold']),
    k: z
      .number()
      .optional()
      .describe('Only for k nearest neighbor (knn) weights'),
    orderOfContiguity: z.number().optional(),
    includeLowerOrder: z.boolean().optional(),
    precisionThreshold: z
      .number()
      .optional()
      .describe(
        'For queen/rook weights: precision threshold for matching coordinates to determine neighboring polygons. Default: 0.'
      ),
    distanceThreshold: z
      .number()
      .optional()
      .describe(
        'Only for distance based weights. It represents the distance threshold used to search nearby neighbors.'
      ),
    useCentroids: z
      .boolean()
      .optional()
      .describe(
        'Whether to use centroids for neighbor calculations. The default value is False.'
      ),
    isMile: z.boolean().optional().describe('Only for distance based weights.'),
    mapBounds: z.array(z.number()).optional(),
  }),
  execute: executeSpatialWeights,
  context: {
    getGeometries: () => {
      throw new Error(
        'getGeometries() of SpatialWeightsTool is not implemented'
      );
    },
  },
};

export type SpatialWeightsTool = typeof spatialWeights;

export type SpatialWeightsFunctionContext = {
  getGeometries: GetGeometries;
};

export type GetWeights = (
  datasetName: string,
  type: 'knn' | 'queen' | 'rook' | 'threshold',
  options: {
    k?: number;
    orderOfContiguity?: number;
    includeLowerOrder?: boolean;
    precisionThreshold?: number;
    distanceThreshold?: number;
    isMile?: boolean;
    useCentroids?: boolean;
  }
) => Promise<{
  weights: number[][];
  weightsMeta: WeightsMeta;
}>;

export type ExecuteSpatialWeightsResult = {
  llmResult: {
    success: boolean;
    weightsId?: string;
    weightsMeta?: WeightsMeta;
    result?: string;
    error?: string;
    instruction?: string;
  };
  additionalData?: {
    [id: string]: unknown;
    weightsId: string;
  };
};

type SpatialWeightsArgs = {
  datasetName: string;
  type: 'knn' | 'queen' | 'rook' | 'threshold';
  k?: number;
  orderOfContiguity?: number;
  includeLowerOrder?: boolean;
  precisionThreshold?: number;
  distanceThreshold?: number;
  isMile?: boolean;
  useCentroids?: boolean;
  mapBounds?: number[];
};

function isSpatialWeightsArgs(args: unknown): args is SpatialWeightsArgs {
  return (
    typeof args === 'object' &&
    args !== null &&
    'datasetName' in args &&
    typeof args.datasetName === 'string' &&
    'type' in args &&
    typeof args.type === 'string' &&
    ['knn', 'queen', 'rook', 'threshold'].includes(args.type)
  );
}

function isSpatialWeightsContext(
  context: unknown
): context is SpatialWeightsFunctionContext {
  return (
    typeof context === 'object' &&
    context !== null &&
    'getGeometries' in context &&
    typeof context.getGeometries === 'function'
  );
}

async function executeSpatialWeights(
  args: z.infer<SpatialWeightsFunctionArgs>,
  options?: {
    toolCallId: string;
    abortSignal?: AbortSignal;
    context?: SpatialWeightsFunctionContext;
  }
): Promise<
  OpenAssistantExecuteFunctionResult<
    SpatialWeightsLlmResult,
    SpatialWeightsAdditionalData
  >
> {
  try {
    if (!isSpatialWeightsArgs(args)) {
      throw new Error('Invalid arguments for spatialWeights tool');
    }

    if (!options?.context || !isSpatialWeightsContext(options.context)) {
      throw new Error('Invalid context for spatialWeights tool');
    }

    const {
      datasetName,
      type,
      k,
      orderOfContiguity,
      includeLowerOrder,
      precisionThreshold,
      distanceThreshold,
      isMile,
      useCentroids,
      mapBounds,
    } = args;
    const { getGeometries } = options.context;
    const geometries = await getGeometries(datasetName);

    if (!geometries) {
      throw new Error(
        `Error: geometries are empty. Please implement the getGeometries() context function.`
      );
    }

    const weightsProps: CreateWeightsProps = {
      weightsType: type,
      k,
      isQueen: type === 'queen',
      distanceThreshold,
      isMile,
      useCentroids,
      precisionThreshold,
      orderOfContiguity,
      includeLowerOrder,
      geometries,
    };

    const id = getWeightsId(datasetName, weightsProps, mapBounds);

    let w: { weightsMeta: WeightsMeta; weights: number[][] } | null = null;

    // check if the weights already exist in the global variable
    const existingWeightData = globalWeightsData[id];
    if (existingWeightData) {
      w = {
        weightsMeta: existingWeightData.weightsMeta,
        weights: existingWeightData.weights,
      };
    } else {
      // create the weights if it does not exist
      const result = await createWeights(weightsProps);
      w = {
        weightsMeta: result.weightsMeta,
        weights: result.weights,
      };
    }
    // set the id to the weights meta
    w.weightsMeta.id = id;

    // cache the weights to the global variable, so that it can be reused across tool calls e.g. lisa, spatial regression
    globalWeightsData[id] = {
      datasetId: datasetName,
      ...w,
    };

    return {
      llmResult: {
        success: true,
        weightsId: id,
        weightsMeta: w.weightsMeta,
        result: `Weights created successfully and the weights are saved using weightsId: ${id}.`,
      },
      additionalData: {
        weightsId: id,
        [id]: {
          type: 'weights',
          content: {
            weights: w.weights,
            weightsMeta: w.weightsMeta,
          },
        },
      },
    };
  } catch (error) {
    return {
      llmResult: {
        success: false,
        error: error instanceof Error ? error.message : String(error),
        instruction:
          'Please explain the error and give a plan to fix the error. Then try again with a different weights type.',
      },
    };
  }
}

export function getWeightsId(
  datasetId: string,
  weightsProps: CreateWeightsProps,
  mapBounds?: number[]
): string {
  const parts = [
    'w', // prefix
    datasetId,
    weightsProps.weightsType,
  ];

  if (
    weightsProps.weightsType === 'queen' ||
    weightsProps.weightsType === 'rook'
  ) {
    parts.push(
      String(weightsProps.orderOfContiguity || 1),
      weightsProps.includeLowerOrder ? 'lower' : '',
      String(weightsProps.precisionThreshold || 0)
    );
  } else if (weightsProps.weightsType === 'knn') {
    parts.push(String(weightsProps.k));
  } else if (weightsProps.weightsType === 'threshold') {
    const distanceThresholdString = weightsProps.distanceThreshold
      ? weightsProps.distanceThreshold.toFixed(1)
      : '0';
    parts.push(distanceThresholdString, weightsProps.isMile ? 'mile' : 'km');
  }

  if (mapBounds) {
    parts.push(mapBounds[0].toString());
  }

  return parts.filter(Boolean).join('-');
}

export function getCachedWeights(
  datasetId: string,
  createWeightsProps: CreateWeightsProps
) {
  const id = getWeightsId(datasetId, createWeightsProps);
  const existingWeightData = globalWeightsData[id];
  if (existingWeightData) {
    return {
      weightsMeta: existingWeightData.weightsMeta,
      weights: existingWeightData.weights,
    };
  }
  return null;
}

export function getCachedWeightsById(weightsId: string) {
  const existingWeightData = globalWeightsData[weightsId];
  if (existingWeightData) {
    return existingWeightData;
  }
  return null;
}

export function isWeightsAdditionalData(
  data: unknown
): data is SpatialWeightsAdditionalData {
  if (typeof data === 'object' && data !== null && 'weightsId' in data) {
    const weightsId = data.weightsId;
    if (typeof weightsId === 'string' && data[weightsId] !== undefined) {
      return true;
    }
  }
  return false;
}