import { vectorTilesetSource } from '@carto/api-client'
import {
  useQuery,
  type UseQueryOptions,
  type UseQueryResult,
} from '@tanstack/react-query'
import { getQueryStats } from './get-query-stats'
import { getTableStats } from './get-table-stats'
import type {
  AllExternalSources,
  GeoJSONSource,
  QuerySource,
  Source,
  TableSource,
  TilesetSource,
} from '../providers/sources/types'
import type { Stats, StatsResponse } from './types'
import type { Feature } from 'geojson'
import { useQueryKey } from '../models/utils/use-query-key'

/**
 * Requests statistics for a specific property of a data source, returning aggregate metrics for numeric attributes or category distributions for string attributes.
 *
 * @returns A React Query result containing `data` (typed as `StringStatsResponse` or `NumberStatsResponse`), `error`, and `isLoading`.
 *
 * @remarks
 * The hook is generic: use `useStats<number>` for numeric attributes (returns avg, sum, min, max, quantiles) or `useStats<string>` for string attributes (returns category frequencies).
 * Supports CARTO query/table/tileset sources and client-side GeoJSON sources.
 *
 * @example
 * ```tsx
 * const { data, error, isLoading } = useStats<number>({
 *   source: mySource,
 *   attribute: 'population',
 *   accessToken: 'YOUR_ACCESS_TOKEN',
 * })
 * ```
 */
export function useStats<Attribute extends string | number>({
  source,
  attribute,
  accessToken,
  params,
  fetcherOptions,
  useQueryProps: { enabled = true, ...useQueryProps } = {},
}: {
  source: Source | undefined
  attribute: Stats['attribute']
  accessToken: Stats['accessToken']
  params?: Stats['params']
  fetcherOptions?: Stats['fetcherOptions']
  useQueryProps?: Omit<
    UseQueryOptions<StatsResponse<Attribute>, Error>,
    'queryKey' | 'queryFn'
  >
}): UseQueryResult<StatsResponse<Attribute>, Error> {
  const queryKey = useQueryKey({
    queryKeyPrefix: ['stats', attribute],
    sourceId: source?.id,
  })

  return useQuery<StatsResponse<Attribute>, Error>({
    queryKey,
    queryFn: () =>
      retrieveStatsBySource<Attribute>({
        source: source!,
        params,
        attribute,
        accessToken: accessToken!,
        fetcherOptions,
      }),
    select(data) {
      return {
        ...data,
        type: data.type.toLowerCase(),
      } as StatsResponse<Attribute>
    },
    enabled: !!source && enabled,
    ...useQueryProps,
  })
}

function retrieveStatsBySource<Attribute extends string | number>({
  source,
  params,
  attribute,
  accessToken,
  fetcherOptions,
}: {
  source: Source
  params?: Stats['params']
  attribute: Stats['attribute']
  accessToken: NonNullable<Stats['accessToken']>
  fetcherOptions?: Stats['fetcherOptions']
}): Promise<StatsResponse<Attribute>> {
  if (source.type.startsWith('query')) {
    const _source = source as Source<AllExternalSources & QuerySource>
    return getQueryStats<Attribute>({
      connection: _source.connectionName,
      baseUrl: _source.apiBaseUrl,
      apiVersion: _source.apiVersion,
      accessToken,
      params,
      query: _source.data,
      queryParameters: _source.queryParameters,
      attribute,
      fetcherOptions,
    })
  }
  if (source.type.startsWith('table')) {
    const _source = source as Source<AllExternalSources & TableSource>
    return getTableStats<Attribute>({
      connection: _source.connectionName,
      baseUrl: _source.apiBaseUrl,
      apiVersion: _source.apiVersion,
      params,
      accessToken,
      table: _source.data,
      attribute,
      fetcherOptions,
    })
  }
  if (source.type.startsWith('tileset')) {
    const _source = source as Source<TilesetSource>
    return getStatsFromTileset<Attribute>({
      source: _source,
      accessToken,
      target: _source.data,
      attribute,
    })
  }

  if (source.type === 'geojson') {
    return Promise.resolve(
      getStatsFromGeoJSON({
        source: source as Source<GeoJSONSource>,
        attribute,
      }) as StatsResponse<Attribute>,
    )
  }

  throw new Error(`Unsupported source type: ${source.type}`)
}

// Catch Stats from tileset

async function getStatsFromTileset<Attribute extends string | number>({
  source,
  accessToken,
  target,
  attribute,
}: {
  source: Source<AllExternalSources & TilesetSource>
  accessToken: string
  target: string
  attribute: string
}) {
  const tileJson = await vectorTilesetSource({
    connectionName: source.connectionName,
    apiBaseUrl: source.apiBaseUrl,
    accessToken: accessToken,
    clientId: source.clientId,
    tableName: target,
  })
  if (!tileJson?.tilestats?.layers?.[0]?.attributes) {
    throw new Error(`${attribute} not found in tileset attributes`)
  }

  const tileStatsAttributes = tileJson.tilestats.layers[0].attributes
  const columnStats = tileStatsAttributes.find((p) => p.attribute === attribute)

  if (!columnStats) {
    throw new Error(`${attribute} not found in tileset attributes`)
  }

  return columnStats as StatsResponse<Attribute>
}

function getStatsFromGeoJSON({
  source,
  attribute,
}: {
  source: Source<GeoJSONSource>
  attribute: string
}) {
  const data = source.data
  // Find the first non-null, non-undefined value for the attribute
  const attrType = data.find((feature) => {
    const value = feature?.properties?.[attribute] as unknown
    return value !== null && value !== undefined
  })?.properties?.[attribute] as unknown

  if (typeof attrType === 'number') {
    return getNumberStats({ data, attribute })
  }

  return getStringStats({ data, attribute })
}

function getNumberStats({
  data,
  attribute,
}: {
  data: Feature[]
  attribute: string
}): StatsResponse<number> {
  const _response = data.reduce(
    (acc, curr) => {
      const value = curr.properties?.[attribute] as number

      acc.sum += value
      acc.max = Math.max(acc.max, value)
      acc.min = Math.min(acc.min, value)
      acc.values.push(value)

      return acc
    },
    {
      max: -Infinity,
      min: Infinity,
      sum: 0,
      values: [],
    } as {
      max: number
      min: number
      sum: number
      values: number[]
    },
  )

  const values = _response.values.sort((a, b) => a - b)

  const quantiles: Record<number, number[]> = {}
  for (let i = 3; i <= 20; i++) {
    quantiles[i] = []
    for (let j = 0; j <= i; j++) {
      const qIndex = Math.floor((j * (values.length - 1)) / i)
      const value = values[qIndex]
      if (value !== undefined) {
        quantiles[i]?.push(value)
      }
    }
  }

  const response = {
    ..._response,
    avg: data.length ? _response.sum / data.length : 0,
    attribute,
    quantiles,
    type: 'number' as const,
  }

  return response
}

function getStringStats({
  data,
  attribute,
}: {
  data: Feature[]
  attribute: string
}): StatsResponse<string> {
  const valuesCount: Record<string, number> = {}
  data.forEach((feature) => {
    const value = feature.properties?.[attribute] as string
    if (value !== undefined) {
      valuesCount[value] ??= 0
      valuesCount[value] += 1
    }
  })

  const categories = Object.entries(valuesCount).map(([value, count]) => ({
    category: value,
    frequency: count,
  }))

  const response = {
    categories,
    type: 'string' as const,
    attribute,
  }

  return response
}
