import path from "path"
import z from "zod"
import { Global } from "../global"
import { JsonStore } from "../util/jsonstore"

export namespace McpAuth {
  export const Tokens = z.object({
    accessToken: z.string(),
    refreshToken: z.string().optional(),
    expiresAt: z.number().optional(),
    scope: z.string().optional(),
  })
  export type Tokens = z.infer<typeof Tokens>

  export const ClientInfo = z.object({
    clientId: z.string(),
    clientSecret: z.string().optional(),
    clientIdIssuedAt: z.number().optional(),
    clientSecretExpiresAt: z.number().optional(),
  })
  export type ClientInfo = z.infer<typeof ClientInfo>

  export const Entry = z.object({
    tokens: Tokens.optional(),
    clientInfo: ClientInfo.optional(),
    codeVerifier: z.string().optional(),
    oauthState: z.string().optional(),
    serverUrl: z.string().optional(), // Track the URL these credentials are for
  })
  export type Entry = z.infer<typeof Entry>

  const filepath = path.join(Global.Path.data, "mcp-auth.json")

  export async function get(mcpName: string): Promise<Entry | undefined> {
    const data = await all()
    return data[mcpName]
  }

  /**
   * Get auth entry and validate it's for the correct URL.
   * Returns undefined if URL has changed (credentials are invalid).
   */
  export async function getForUrl(mcpName: string, serverUrl: string): Promise<Entry | undefined> {
    const entry = await get(mcpName)
    if (!entry) return undefined

    // If no serverUrl is stored, this is from an old version - consider it invalid
    if (!entry.serverUrl) return undefined

    // If URL has changed, credentials are invalid
    if (entry.serverUrl !== serverUrl) return undefined

    return entry
  }

  export async function all(): Promise<Record<string, Entry>> {
    return (await JsonStore.read(filepath)) as Record<string, Entry>
  }

  /** Atomic read-modify-write of one server's entry. Serialized behind the
   *  store lock so concurrent updates can't drop each other's writes. */
  async function update(mcpName: string, fn: (entry: Entry) => void, serverUrl?: string): Promise<void> {
    await JsonStore.update(filepath, (data) => {
      const entry = (data[mcpName] ?? {}) as Entry
      fn(entry)
      // Always update serverUrl if provided
      if (serverUrl) entry.serverUrl = serverUrl
      data[mcpName] = entry
    })
  }

  export async function set(mcpName: string, entry: Entry, serverUrl?: string): Promise<void> {
    // Always update serverUrl if provided
    if (serverUrl) entry.serverUrl = serverUrl
    await JsonStore.update(filepath, (data) => ({ ...data, [mcpName]: entry }))
  }

  export async function remove(mcpName: string): Promise<void> {
    await JsonStore.update(filepath, (data) => {
      delete data[mcpName]
    })
  }

  export async function updateTokens(mcpName: string, tokens: Tokens, serverUrl?: string): Promise<void> {
    await update(
      mcpName,
      (entry) => {
        entry.tokens = tokens
      },
      serverUrl,
    )
  }

  export async function updateClientInfo(mcpName: string, clientInfo: ClientInfo, serverUrl?: string): Promise<void> {
    await update(
      mcpName,
      (entry) => {
        entry.clientInfo = clientInfo
      },
      serverUrl,
    )
  }

  export async function updateCodeVerifier(mcpName: string, codeVerifier: string): Promise<void> {
    await update(mcpName, (entry) => {
      entry.codeVerifier = codeVerifier
    })
  }

  export async function clearCodeVerifier(mcpName: string): Promise<void> {
    await JsonStore.update(filepath, (data) => {
      const entry = data[mcpName] as Entry | undefined
      if (entry) delete entry.codeVerifier
    })
  }

  export async function updateOAuthState(mcpName: string, oauthState: string): Promise<void> {
    await update(mcpName, (entry) => {
      entry.oauthState = oauthState
    })
  }

  export async function getOAuthState(mcpName: string): Promise<string | undefined> {
    const entry = await get(mcpName)
    return entry?.oauthState
  }

  export async function clearOAuthState(mcpName: string): Promise<void> {
    await JsonStore.update(filepath, (data) => {
      const entry = data[mcpName] as Entry | undefined
      if (entry) delete entry.oauthState
    })
  }

  /**
   * Check if stored tokens are expired.
   * Returns null if no tokens exist, false if no expiry or not expired, true if expired.
   */
  export async function isTokenExpired(mcpName: string): Promise<boolean | null> {
    const entry = await get(mcpName)
    if (!entry?.tokens) return null
    if (!entry.tokens.expiresAt) return false
    return entry.tokens.expiresAt < Date.now() / 1000
  }
}