import { readFileSync } from "node:fs"
import { join } from "node:path"
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"
import { getAgentDir } from "@earendil-works/pi-coding-agent"
import type { AutocompleteItem } from "@earendil-works/pi-tui"

type ParamName = 'temperature' | 'top_p' | 'top_k' | 'min_p' | 'frequency_penalty' | 'presence_penalty' | 'repetition_penalty'

const DEFAULTS: Record<ParamName, number | undefined> = {
  temperature: 1.0,
  top_p: 0.95,
  top_k: undefined,
  min_p: undefined,
  frequency_penalty: undefined,
  presence_penalty: undefined,
  repetition_penalty: undefined,
}

const PARAM_NAMES = Object.keys(DEFAULTS) as ParamName[]

function loadModelParams(modelId: string): Record<ParamName, number | undefined> {
  const result: Record<ParamName, number | undefined> = { ...DEFAULTS }
  try {
    const modelsPath = join(getAgentDir(), "models.json")
    const raw = JSON.parse(readFileSync(modelsPath, "utf8"))
    const providers = raw.providers as Record<string, { models?: Array<{ id: string; parameters?: Record<string, number> }> }> | undefined
    if (providers) {
      for (const provider of Object.values(providers)) {
        const match = provider.models?.find((m) => m.id === modelId)
        if (match?.parameters) {
          for (const key of PARAM_NAMES) {
            result[key] = key in match.parameters ? match.parameters[key] : undefined
          }
          break
        }
      }
    }
  } catch {}
  return result
}

export default function init(pi: ExtensionAPI) {
  const params: Record<ParamName, number | undefined> = { ...DEFAULTS }

  function setParam(key: string, value: string): string {
    if (!(key in params)) return `Unknown parameter "${key}". Valid parameters: ${PARAM_NAMES.join(', ')}`
    const num = parseFloat(value)
    if (isNaN(num)) return `"${value}" is not a number.`
    params[key as ParamName] = num
    return `Set ${key} = ${num}`
  }

  function unsetParam(key: string): string {
    if (!(key in params)) return `Unknown parameter "${key}". Valid parameters: ${PARAM_NAMES.join(', ')}`
    params[key as ParamName] = undefined
    return `Unset ${key} (${key} will no longer be sent)`
  }

  function formatStatus(): string {
    const active = Object.entries(params)
      .filter(([_, v]) => v != null)
      .map(([k, v]) => `- \`${k}\`: ${v}`)
    return active.length ? 'Active parameters:\n' + active.join('\n') : 'No custom parameters are sent.'
  }

  pi.on('session_start', (_event, ctx) => {
    if (ctx.model) {
      Object.assign(params, loadModelParams(ctx.model.id))
    }
  })

  pi.on('model_select', (event) => {
    Object.assign(params, loadModelParams(event.model.id))
  })

  pi.on('before_provider_request', (event: any) => {
    const payload = event.payload
    if (payload && typeof payload === 'object') {
      const active = Object.fromEntries(
        Object.entries(params).filter(([_, v]) => v != null)
      )
      Object.assign(payload, active)
    }
    return payload
  })

  pi.registerCommand('params', {
    description: `View/set LLM parameters (${PARAM_NAMES.join(', ')})`,
    getArgumentCompletions: (prefix: string): AutocompleteItem[] | null => {
      const parts = prefix.split(/\s+/)

      if (parts.length <= 1) {
        const subs = ['status', 'set', 'unset']
        const filtered = subs.filter((s) => s.startsWith(parts[0] || ''))
        return filtered.length > 0 ? filtered.map((s) => ({ value: s, label: s })) : null
      }

      if ((parts[0] === 'set' || parts[0] === 'unset') && parts.length === 2) {
        const filtered = PARAM_NAMES.filter((n) => n.startsWith(parts[1] || ''))
        return filtered.length > 0 ? filtered.map((n) => ({ value: `${parts[0]} ${n}`, label: n })) : null
      }

      return null
    },
    handler: async (args, ctx) => {
      const parts = args.trim().split(/\s+/)
      const sub = parts[0]?.toLowerCase()

      if (!sub || sub === 'status') {
        ctx.ui.notify(formatStatus(), 'info')
        return
      }

      if (sub === 'set') {
        const key = parts[1]
        const value = parts[2]
        if (!key || !value) {
          ctx.ui.notify('Usage: /params set <key> <value>', 'info')
          return
        }
        ctx.ui.notify(setParam(key, value), 'info')
        return
      }

      if (sub === 'unset') {
        const key = parts[1]
        if (!key) {
          ctx.ui.notify('Usage: /params unset <key>', 'info')
          return
        }
        ctx.ui.notify(unsetParam(key), 'info')
        return
      }

      ctx.ui.notify('Usage: `/params [status|set <key> <value>|unset <key>]`', 'info')
    },
  })
}