import { Client } from "./client.js"
import { Identity } from "./identity.js"
import { TelemetryEvent } from "./events.js"

export interface TelemetryProperties {
  appName: string
  appVersion: string
  platform: string
  editorName?: string
  vscodeVersion?: string
}

export type ReviewCommand = "review" | "local-review" | "local-review-uncommitted"

export interface IndexingTelemetryProperties extends Record<string, unknown> {
  source: "scan" | "watcher"
  provider: string
  vectorStore: "lancedb" | "qdrant"
  modelId?: string
  trigger?: "background" | "manual"
  mode?: "full" | "incremental"
}

export interface IndexingCompletedTelemetryProperties extends IndexingTelemetryProperties {
  trigger: "background" | "manual"
  mode: "full" | "incremental"
  filesIndexed: number
  filesDiscovered: number
  totalBlocks: number
  batchErrors: number
}

export interface IndexingFileCountTelemetryProperties extends IndexingTelemetryProperties {
  mode: "full" | "incremental"
  discovered: number
  candidate: number
}

export interface IndexingRetryTelemetryProperties extends IndexingTelemetryProperties {
  mode: "full" | "incremental"
  attempt: number
  maxRetries: number
  batchSize: number
  error: string
}

export interface IndexingErrorTelemetryProperties extends IndexingTelemetryProperties {
  location: string
  error: string
  retryCount?: number
  maxRetries?: number
}

export namespace Telemetry {
  let initialized = false
  let startTime = 0
  let props: TelemetryProperties = {
    appName: "kilo-cli",
    appVersion: "unknown",
    platform: process.platform,
  }

  export async function init(options: { dataPath: string; version: string; enabled: boolean }): Promise<void> {
    if (initialized) return

    Identity.setDataPath(options.dataPath)
    props.appVersion = options.version

    const app = process.env.KILO_APP_NAME
    if (app) props.appName = app
    const editor = process.env.KILO_EDITOR_NAME
    if (editor) props.editorName = editor
    const platform = process.env.KILO_PLATFORM
    if (platform) props.platform = platform
    const version = process.env.KILO_APP_VERSION
    if (version) props.appVersion = version
    const vscodeVersion = process.env.KILO_VSCODE_VERSION
    if (vscodeVersion) props.vscodeVersion = vscodeVersion

    Client.init()

    const level = process.env.KILO_TELEMETRY_LEVEL
    const enabled = level ? level === "all" : options.enabled
    Client.setEnabled(enabled)

    await Identity.getMachineId()

    initialized = true
    startTime = Date.now()
  }

  export function setEnabled(value: boolean) {
    Client.setEnabled(value)
  }

  export function isEnabled(): boolean {
    return Client.isEnabled()
  }

  export async function updateIdentity(token: string | null, accountId?: string): Promise<void> {
    const previousId = Identity.getDistinctId()
    await Identity.updateFromKiloAuth(token, accountId)

    const email = Identity.getUserId()
    if (email && previousId && email !== previousId) {
      // Identify the user with their email and properties
      Client.identify(email, {
        ...(accountId && { kilocodeOrganizationId: accountId }),
        appName: props.appName,
        appVersion: props.appVersion,
        platform: props.platform,
      })

      // Link the anonymous machineId to the authenticated email
      Client.alias(email, previousId)
    }
  }

  export function track(event: TelemetryEvent, properties?: Record<string, unknown>) {
    Client.capture(event, { ...props, ...properties })
  }

  // CLI Lifecycle
  export function trackCliStart() {
    track(TelemetryEvent.CLI_START)
  }

  export function trackCliExit(exitCode?: number) {
    track(TelemetryEvent.CLI_EXIT, {
      duration: Date.now() - startTime,
      exitCode,
    })
  }

  // Sessions
  export function trackSessionStart(sessionId: string, model?: string, provider?: string) {
    track(TelemetryEvent.SESSION_START, { sessionId, model, provider })
  }

  export function trackSessionEnd(
    sessionId: string,
    stats: {
      messageCount?: number
      inputTokens?: number
      outputTokens?: number
      duration?: number
    },
  ) {
    track(TelemetryEvent.SESSION_END, { sessionId, ...stats })
  }

  export function trackSessionMessage(sessionId: string, source: "user" | "assistant") {
    track(TelemetryEvent.SESSION_MESSAGE, { sessionId, source })
  }

  // LLM
  export function trackLlmCompletion(properties: {
    taskId?: string
    mode?: "review"
    feature?: "code_reviews"
    command?: ReviewCommand
    tool?: "suggest"
    apiProvider: string
    modelId: string
    inputTokens?: number
    outputTokens?: number
    cacheReadTokens?: number
    cacheWriteTokens?: number
    cost?: number
    completionTime?: number
    duration?: number
  }) {
    track(TelemetryEvent.LLM_COMPLETION, properties)
  }

  // Features
  export function trackCommandUsed(command: string) {
    track(TelemetryEvent.COMMAND_USED, { command })
  }

  export function trackToolUsed(tool: string, sessionId?: string) {
    track(TelemetryEvent.TOOL_USED, { tool, sessionId })
  }

  export function trackAgentUsed(agent: string, sessionId?: string) {
    track(TelemetryEvent.AGENT_USED, { agent, sessionId })
  }

  export function trackPlanFollowup(sessionId: string, choice: "new_session" | "continue" | "custom" | "dismissed") {
    track(TelemetryEvent.PLAN_FOLLOWUP, { sessionId, choice })
  }

  export function trackSuggestionAccepted(properties: {
    sessionId: string
    requestId: string
    index: number
    tool: "suggest"
    command: ReviewCommand
    actionCount?: number
  }) {
    track(TelemetryEvent.SUGGESTION_ACCEPTED, properties)
  }

  export function trackSuggestionShown(properties: {
    sessionId: string
    requestId: string
    index: number
    tool: "suggest"
    command: ReviewCommand
    actionCount?: number
  }) {
    track(TelemetryEvent.SUGGESTION_SHOWN, properties)
  }

  export function trackIndexingStarted(properties: IndexingTelemetryProperties) {
    track(TelemetryEvent.INDEXING_STARTED, properties)
  }

  export function trackIndexingCompleted(properties: IndexingCompletedTelemetryProperties) {
    track(TelemetryEvent.INDEXING_COMPLETED, properties)
  }

  export function trackIndexingFileCount(properties: IndexingFileCountTelemetryProperties) {
    track(TelemetryEvent.INDEXING_FILE_COUNT, properties)
  }

  export function trackIndexingBatchRetry(properties: IndexingRetryTelemetryProperties) {
    track(TelemetryEvent.INDEXING_BATCH_RETRY, properties)
  }

  export function trackIndexingError(properties: IndexingErrorTelemetryProperties) {
    track(TelemetryEvent.INDEXING_ERROR, properties)
  }

  // Share
  export function trackShareCreated(sessionId: string) {
    track(TelemetryEvent.SHARE_CREATED, { sessionId })
  }

  export function trackShareDeleted(sessionId: string) {
    track(TelemetryEvent.SHARE_DELETED, { sessionId })
  }

  // MCP
  export function trackMcpServerConnected(server: string) {
    track(TelemetryEvent.MCP_SERVER_CONNECTED, { server })
  }

  export function trackMcpServerError(server: string, error?: string) {
    track(TelemetryEvent.MCP_SERVER_ERROR, { server, error })
  }

  // Remote
  export function trackRemoteConnectionOpened() {
    track(TelemetryEvent.REMOTE_CONNECTION_OPENED)
  }

  // Auth
  export function trackAuthSuccess(provider: string) {
    track(TelemetryEvent.AUTH_SUCCESS, { provider })
  }

  export function trackAuthLogout(provider: string) {
    track(TelemetryEvent.AUTH_LOGOUT, { provider })
  }

  // Errors
  export function trackError(error: string, context?: string) {
    track(TelemetryEvent.ERROR, { error, context })
  }

  // Feedback
  export interface FeedbackProperties extends Record<string, unknown> {
    providerID: string
    modelID: string
    variant?: string
    rating: "up" | "down" | "cleared"
    previousRating?: "up" | "down"
    sessionID?: string
    messageID?: string
    parentMessageID?: string
  }

  export function trackFeedback(props: FeedbackProperties) {
    track(TelemetryEvent.FEEDBACK_SUBMITTED, props)
  }

  export async function shutdown(timeoutMs?: number): Promise<void> {
    await Client.shutdown(timeoutMs)
  }
}
