dotfiles

# Kimi Code / Moonshot AI integration.
# Set KIMI_API_KEY in secrets.nuon.
# Override base URL with KIMI_API_BASE (default: https://api.kimi.com/coding/v1).

export def main [] {
    print "Kimi Code / Moonshot AI integration"
    print "  kimi usage           — check subscription quota"
    print "  kimi chat <message>  — send a chat completion"
    print "  kimi models          — list available models"
    print ""
    print "Run a subcommand with --help for details."
}

# ── Usage ─────────────────────────────────────────────────────────────────────

# `kimi usage` — check Kimi Code subscription quota and usage.
export def usage [] {
    let token = get-token
    let base = get-base
    let hs = {
        Authorization: $'Bearer ($token)'
        Accept: "application/json"
    }

    let res = try {
        http get --headers $hs $"($base)/usages"
    } catch { |e|
        error make {
            msg: "Failed to fetch usage"
            label: { text: "Check KIMI_API_KEY and network connection" }
        }
    }

    let c = (ansi c)
    let g = (ansi g)
    let y = (ansi y)
    let r = (ansi r)
    let d = (ansi d)
    let e = (ansi rst)

    print $"($c)KIMI CODE USAGE($e)"
    print ""

    # ── weekly / main usage ──
    let usage = $res | get --optional usage | default {}
    let limit = $usage | get --optional limit | default "0"
    let used = $usage | get --optional used | default "0"
    let remain = $usage | get --optional remaining | default "0"
    let reset = $usage | get --optional resetTime | default ""

    if ($limit | into int) > 0 {
        let pct = ($remain | into int) * 100 // ($limit | into int)
        let bar = bar $pct
        let color = if $pct <= 10 { $r } else if $pct <= 30 { $y } else { $g }
        let reset_hint = if $reset != "" { $" ($d)resets ($reset)($e)" } else { "" }
        print $"  ($c)Weekly Quota($e)"
        print $"    ($bar) ($color)($remain)/($limit) ($pct)% left($e)($reset_hint)"
        print ""
    }

    # ── total quota ──
    let total_quota = $res | get --optional totalQuota | default {}
    let tq_limit = $total_quota | get --optional limit | default "0"
    let tq_remain = $total_quota | get --optional remaining | default "0"
    if ($tq_limit | into int) > 0 {
        print $"  ($c)Total Quota($e)  ($tq_remain)/($tq_limit) remaining"
        print ""
    }

    # ── windowed limits ──
    let limits = $res | get --optional limits | default []
    if ($limits | length) > 0 {
        print $"  ($c)Rate Limits($e)"
        for item in $limits {
            let window = $item | get --optional window | default {}
            let detail = $item | get --optional detail | default {}
            let w_dur = $window | get --optional duration | default 0
            let w_unit = $window | get --optional timeUnit | default ""
            let d_limit = $detail | get --optional limit | default "0"
            let d_used = $detail | get --optional used | default "0"
            let d_remain = $detail | get --optional remaining | default "0"

            let label = if ($w_unit | str contains "MINUTE") and $w_dur >= 60 and ($w_dur mod 60) == 0 {
                $"($w_dur / 60)h limit"
            } else if ($w_unit | str contains "MINUTE") {
                $"($w_dur)m limit"
            } else if ($w_unit | str contains "HOUR") {
                $"($w_dur)h limit"
            } else if ($w_unit | str contains "DAY") {
                $"($w_dur)d limit"
            } else {
                "Rate limit"
            }

            let pct = if ($d_limit | into int) > 0 {
                ($d_remain | into int) * 100 // ($d_limit | into int)
            } else { 0 }
            let bar = bar $pct
            print $"    ($label)  ($bar) ($d_remain)/($d_limit)"
        }
        print ""
    }

    # ── parallel / concurrency ──
    let parallel = $res | get --optional parallel | default {}
    let p_limit = $parallel | get --optional limit | default "0"
    if ($p_limit | into int) > 0 {
        print $"  ($c)Concurrency($e)  max ($p_limit) parallel requests"
        print ""
    }
}

def bar [pct: int] {
    let n = ($pct / 10 | math round | into int)
    let n = if $n > 10 { 10 } else if $n < 0 { 0 } else { $n }
    let g = (ansi g)
    let d = (ansi d)
    let y = (ansi y)
    let r = (ansi r)
    let e = (ansi rst)
    let color = if $pct < 50 { $g } else if $pct < 80 { $y } else { $r }
    let filled = 1..$n | each { $"($g)▓" } | str join
    let empty = 1..(10 - $n) | each { $"($d)░" } | str join
    $"($filled)($empty)($e)"
}

# ── Private helpers ───────────────────────────────────────────────────────────

def get-token [] {
    let token = $env | get --optional KIMI_API_KEY
    if ($token | is-empty) {
        error make {
            msg: "KIMI_API_KEY is not set"
            label: { text: "Set KIMI_API_KEY in secrets.nuon" }
        }
    }
    $token
}

def get-base [] {
    $env | get --optional KIMI_API_BASE | default "https://api.kimi.com/coding/v1"
}

def headers [] {
    {
        Authorization: $'Bearer (get-token)'
        Content-Type: "application/json"
        Accept: "application/json"
    }
}

# ── Chat ──────────────────────────────────────────────────────────────────────

# `kimi chat` — send a chat completion to kimi-for-coding.
#
# # Parameters
# `message`  # The user message (required)
# `--system` # Optional system prompt
export def chat [
    message: string
    --system (-s): string = ""
] {
    if ($message | str trim | is-empty) {
        error make {
            msg: "Message cannot be empty"
            label: { text: "Provide a non-empty message" }
        }
    }

    let base = get-base
    let hs = headers

    mut messages = []
    if ($system | str trim | is-not-empty) {
        $messages = $messages | append { role: "system", content: $system }
    }
    $messages = $messages | append { role: "user", content: $message }

    let payload = {
        model: "kimi-for-coding"
        messages: $messages
    }

    let resp = try {
        http post --headers $hs $"($base)/chat/completions" ($payload | to json)
    } catch { |e|
        let raw = $e.raw | default ""
        if ($raw | str contains "401") {
            error make {
                msg: "Authentication failed (401)"
                label: { text: "Check KIMI_API_KEY is correct and active" }
            }
        } else if ($raw | str contains "403") {
            error make {
                msg: "Access forbidden (403)"
                label: { text: "Kimi Code API key may be restricted. Ensure subscription includes API access and key is active at https://www.kimi.com/code/console" }
            }
        } else if ($raw | str contains "429") {
            error make {
                msg: "Rate limited (429)"
                label: { text: "Too many requests — wait and retry" }
            }
        }
        error make {
            msg: "Chat request failed"
            label: { text: ($e.msg | default "Check network and API base URL") }
        }
    }

    let choice = $resp | get --optional choices | default [] | get 0 | default {}
    let content = $choice | get --optional message.content | default ""
    let finish = $choice | get --optional finish_reason | default ""

    print $content

    if $finish != "stop" and ($finish | is-not-empty) {
        print $"(ansi d)[finish_reason: ($finish)](ansi rst)"
    }
}

# ── Models ────────────────────────────────────────────────────────────────────

# `kimi models` — list models available on the current endpoint.
export def models [] {
    let base = get-base
    let hs = headers

    let res = try {
        http get --headers $hs $"($base)/models"
    } catch { |e|
        error make {
            msg: "Failed to list models"
            label: { text: "Check KIMI_API_KEY and KIMI_API_BASE" }
        }
    }

    $res.data | each { |m|
        {
            id: $m.id
            display_name: ($m | get --optional display_name | default "")
            context: ($m | get --optional context_length | default 0)
        }
    } | table -e
}