// opencode-fixes — repairs common self-hosted inference (llama.cpp / llama-swap / vLLM) quirks in OpenCode v2.
//
// 1. bool-args: some servers return "True"/"FALSE" strings for boolean tool arguments.
//    This plugin coerces strings to real booleans at provably-boolean schema positions.
// 2. toolcall-leak: servers intermittently fail to parse a model's tool call and return it
//    as ordinary assistant text with finish_reason="stop" — OpenCode then treats the turn as
//    a final answer and the session stalls. This plugin detects such turns and asks the model
//    to re-issue the action as a real tool call.
//
// Install: copy this file to ~/.config/opencode/plugins/opencode-fixes.ts (auto-loaded),
// or register a path in opencode.json: "plugin": ["/abs/path/to/opencode-fixes/index.ts"].
// Log: ~/.local/state/opencode/opencode-fixes.log (override with OPENCODE_FIXES_LOG).

import { Plugin } from "@opencode/plugin"
import fs from "node:fs"
import os from "node:os"
import path from "node:path"

const LOG_PATH = process.env.OPENCODE_FIXES_LOG ?? path.join(os.homedir(), ".local", "state", "opencode", "opencode-fixes.log")

function log(line: string) {
  try {
    fs.mkdirSync(path.dirname(LOG_PATH), { recursive: true })
    fs.appendFileSync(LOG_PATH, `${new Date().toISOString()} ${line}\n`)
  } catch {}
}

// ============================================================
// 1. bool-args repair
// ============================================================

const MAX_DEPTH = 6

function lowerBoolean(value: unknown): unknown {
  if (typeof value !== "string") return value
  const lower = value.toLowerCase()
  if (lower === "true") return true
  if (lower === "false") return false
  return value
}

function looksLikeAst(node: any): boolean {
  return !!node && typeof node === "object" && typeof node._tag === "string" && "~effect/Schema" in node
}

// --- Effect Schema AST walker (only rewrites strings at provably-boolean positions) ---

function astAcceptsBoolean(node: any, depth = 0): boolean {
  if (!node || typeof node !== "object" || depth > MAX_DEPTH) return false
  switch (node._tag) {
    case "Boolean":
      return true
    case "Literal":
      return typeof node.literal === "boolean"
    case "Refinement":
      return astAcceptsBoolean(node.from, depth + 1)
    case "Transformation":
      return astAcceptsBoolean(node.to, depth + 1) || astAcceptsBoolean(node.from, depth + 1)
    case "Union":
      return Array.isArray(node.types) && node.types.some((t: any) => astAcceptsBoolean(t, depth + 1))
    default:
      return false
  }
}

function astObjectMembers(node: any, depth = 0): Record<string, any> | undefined {
  if (!node || typeof node !== "object" || depth > MAX_DEPTH) return undefined
  if (Array.isArray(node.propertySignatures)) {
    const result: Record<string, any> = {}
    for (const signature of node.propertySignatures) {
      if (signature && typeof signature.name === "string" && signature.type) result[signature.name] = signature.type
    }
    return result
  }
  if (node._tag === "Refinement") return astObjectMembers(node.from, depth + 1)
  if (node._tag === "Transformation") return astObjectMembers(node.to, depth + 1) ?? astObjectMembers(node.from, depth + 1)
  if (node._tag === "Union" && Array.isArray(node.types)) {
    let merged: Record<string, any> | undefined
    for (const member of node.types) {
      const members = astObjectMembers(member, depth + 1)
      if (!members) continue
      merged = merged ? { ...merged, ...members } : members
    }
    return merged
  }
  return undefined
}

function coerceAst(value: unknown, node: any, depth = 0): unknown {
  if (!node || typeof node !== "object" || depth > MAX_DEPTH) return value

  if (typeof value === "string" && astAcceptsBoolean(node)) {
    const fixed = lowerBoolean(value)
    if (fixed !== value) return fixed
  }

  if (value !== null && typeof value === "object" && !Array.isArray(value)) {
    const members = astObjectMembers(node)
    if (!members) return value
    const source = value as Record<string, unknown>
    let result: Record<string, unknown> | undefined
    for (const [key, current] of Object.entries(source)) {
      const member = members[key]
      if (!member) continue
      const fixed = coerceAst(current, member, depth + 1)
      if (fixed !== current) {
        if (!result) result = { ...source }
        result[key] = fixed
      }
    }
    return result ?? value
  }

  if (Array.isArray(value)) {
    const itemNodes: any[] = []
    if (Array.isArray(node.elements)) itemNodes.push(...node.elements)
    if (node.rest) itemNodes.push(node.rest)
    if (node.type) itemNodes.push(node.type)
    if (itemNodes.length === 0) return value
    let result: unknown[] | undefined
    value.forEach((item, index) => {
      const member = itemNodes[index] ?? itemNodes[itemNodes.length - 1]
      const fixed = coerceAst(item, member, depth + 1)
      if (fixed !== item) {
        if (!result) result = [...value]
        result[index] = fixed
      }
    })
    return result ?? value
  }

  return value
}

// --- Plain JSON Schema walker (plugin/MCP tools registered with JSON schemas) ---

function plainJsonSchema(candidate: any): any {
  if (!candidate || typeof candidate !== "object" || typeof candidate === "function") return undefined
  const serialized = (() => {
    try {
      return JSON.parse(JSON.stringify(candidate))
    } catch {
      return undefined
    }
  })()
  if (serialized && typeof serialized === "object" && (serialized.type === "object" || "properties" in serialized || "$ref" in serialized)) {
    return serialized
  }
  return undefined
}

function booleanBranch(schema: any): boolean {
  if (!schema || typeof schema !== "object") return false
  if (schema.type === "boolean") return true
  if (Array.isArray(schema.type) && schema.type.includes("boolean")) return true
  const branches = Array.isArray(schema.anyOf) ? schema.anyOf : Array.isArray(schema.oneOf) ? schema.oneOf : []
  return branches.length > 0 && branches.every((b: any) => booleanBranch(b))
}

function resolveRef(ref: string, root: any): any {
  if (!/^#\/(\$defs|definitions)\/[^/]+$/.test(ref) || !root) return undefined
  const store = ref.startsWith("#/$defs/") ? root.$defs : root.definitions
  if (!store || typeof store !== "object") return undefined
  const name = ref.slice(ref.lastIndexOf("/") + 1).replaceAll("~1", "/").replaceAll("~0", "~")
  return store[name]
}

function coerceJson(value: unknown, schema: any, root: any = schema, depth = 0): unknown {
  if (!schema || typeof schema !== "object" || depth > MAX_DEPTH) return value

  if (typeof schema.$ref === "string") {
    const target = resolveRef(schema.$ref, root)
    return target ? coerceJson(value, target, root, depth + 1) : value
  }

  if (schema.type === undefined && (Array.isArray(schema.anyOf) || Array.isArray(schema.oneOf))) {
    const branches = (Array.isArray(schema.anyOf) ? schema.anyOf : schema.oneOf) ?? []
    if (booleanBranch(schema) && typeof value === "string") {
      const fixed = lowerBoolean(value)
      if (fixed !== value) return fixed
    }
    for (const branch of branches) {
      const fixed = coerceJson(value, branch, root, depth + 1)
      if (fixed !== value) return fixed
    }
    return value
  }

  if (Array.isArray(schema.type)) {
    if (schema.type.includes(typeof value) || (value === null && schema.type.includes("null"))) return value
    if (typeof value === "string" && schema.type.includes("boolean")) {
      const fixed = lowerBoolean(value)
      if (fixed !== value) return fixed
    }
    return value
  }

  if (schema.type === "boolean") return lowerBoolean(value)

  if (schema.type === "object" && value !== null && typeof value === "object" && !Array.isArray(value)) {
    const source = value as Record<string, unknown>
    const properties = schema.properties && typeof schema.properties === "object" ? schema.properties : {}
    let result: Record<string, unknown> | undefined
    for (const [key, current] of Object.entries(source)) {
      const property = key in properties ? properties[key] : schema.additionalProperties
      if (!property || typeof property !== "object") continue
      const fixed = coerceJson(current, property, root, depth + 1)
      if (fixed !== current) {
        if (!result) result = { ...source }
        result[key] = fixed
      }
    }
    return result ?? value
  }

  if (schema.type === "array" && Array.isArray(value)) {
    const items = schema.items
    if (!items || typeof items !== "object") return value
    let result: unknown[] | undefined
    value.forEach((item, index) => {
      const member = Array.isArray(items) ? items[index] ?? schema.additionalItems : items
      if (!member || typeof member !== "object") return
      const fixed = coerceJson(item, member, root, depth + 1)
      if (fixed !== item) {
        if (!result) result = [...value]
        result[index] = fixed
      }
    })
    return result ?? value
  }

  return value
}

// --- Schema resolution per tool ---

type ToolSchema = { kind: "ast"; node: any } | { kind: "json"; schema: any }

const schemaCache = new Map<string, ToolSchema | undefined>()

function resolveToolSchema(tool: any): ToolSchema | undefined {
  if (!tool) return undefined
  for (const key of ["input", "inputSchema", "parameters", "args"]) {
    const candidate = tool[key]
    if (looksLikeAst(candidate?.ast)) return { kind: "ast", node: candidate.ast }
    const plain = plainJsonSchema(candidate)
    if (plain) return { kind: "json", schema: plain }
  }
  return undefined
}

function resolveCached(toolId: string, tool: any): ToolSchema | undefined {
  if (schemaCache.has(toolId)) return schemaCache.get(toolId)
  const resolved = resolveToolSchema(tool)
  schemaCache.set(toolId, resolved)
  return resolved
}

function coerceInput(input: unknown, resolved: ToolSchema): unknown {
  return resolved.kind === "ast" ? coerceAst(input, resolved.node) : coerceJson(input, resolved.schema)
}

function boolSelfTest(resolved: ToolSchema) {
  const cases: Array<[string, unknown, unknown]> = [
    ["True", { command: "echo hi", background: "True" }, { command: "echo hi", background: true }],
    ["FALSE", { command: "echo hi", background: "FALSE" }, { command: "echo hi", background: false }],
    ["lowercase-normalized", { command: "echo hi", background: "true" }, { command: "echo hi", background: true }],
    ["string-field-untouched", { command: "True", background: false }, { command: "True", background: false }],
  ]
  for (const [name, input, expected] of cases) {
    const actual = coerceInput(structuredClone(input), resolved)
    const pass = JSON.stringify(actual) === JSON.stringify(expected)
    log(`bool self-test ${name}: ${pass ? "PASS" : `FAIL got=${JSON.stringify(actual)}`}`)
  }
}

// ============================================================
// 2. tool-call leak retry
// ============================================================

// Length is not evidence of legitimacy: detection is structural (the tags must anchor the
// message), and a leaked shell call with a heredoc body is routinely >1kB — a hard 1200-char
// cap silently disabled the retry in production (missed live: 1604 chars). Cap is opt-in now.
const LEAK_MAX_LEN = Number(process.env.OPENCODE_FIXES_LEAK_MAX_LEN ?? 0) // 0 = no cap
const LEAK_MAX_CONSECUTIVE = Number(process.env.OPENCODE_FIXES_LEAK_MAX_CONSECUTIVE ?? 3)
const LEAK_RETRY_TEXT = "System notice: your previous message contained a tool call rendered as plain text."

// Fenced code blocks are stripped instead of acting as a veto: a legal "the format looks like
// this" example loses its tags and stays a clean answer, while a real leak only loses its
// command body and keeps the anchoring tags (the old blanket `includes("```")` veto missed those).
function stripFences(text: string): string {
  return text.replace(/```[^\n]*\n[\s\S]*?```/g, "")
}

function leakedToolCallText(text: string): boolean {
  const t = stripFences(text).trim()
  if (!t) return false
  if (LEAK_MAX_LEN > 0 && t.length > LEAK_MAX_LEN) return false
  return t.startsWith("<tool_call>") || t.endsWith("</tool_call>")
}

type LeakState = { consecutive: number; lastMessageID?: string; inFlight: boolean }
const leakStates = new Map<string, LeakState>()

function withTimeout<T>(promise: Promise<T>, ms: number, label: string): Promise<T> {
  return Promise.race([
    promise,
    new Promise<T>((_, reject) => setTimeout(() => reject(new Error(`${label} timeout ${ms}ms`)), ms)),
  ])
}

async function handleTurnEnd(ctx: any, sessionID: string) {
  const st = leakStates.get(sessionID) ?? { consecutive: 0, inFlight: false }
  if (st.inFlight) return
  st.inFlight = true
  try {
    const context: any = await withTimeout(ctx.session.context({ sessionID }), 5000, "context")
    const messages: any[] = Array.isArray(context) ? context : []
    const lastAssistant = [...messages].reverse().find((m: any) => m?.type === "assistant")
    if (!lastAssistant) return
    if (st.lastMessageID === lastAssistant.id) return
    st.lastMessageID = lastAssistant.id

    const content: any[] = Array.isArray(lastAssistant.content) ? lastAssistant.content : []
    const hasTool = content.some((p: any) => String(p?.type ?? "").includes("tool"))
    const text = content
      .filter((p: any) => p?.type === "text")
      .map((p: any) => String(p.text ?? ""))
      .join("\n")
    const isLeak = lastAssistant.finish === "stop" && !hasTool && leakedToolCallText(text)

    if (!isLeak) {
      if (st.consecutive > 0) log(`leak: clean turn, consecutive reset (${sessionID})`)
      st.consecutive = 0
      leakStates.set(sessionID, st)
      return
    }

    // Single-winner dedup across plugin instances (hot reload can keep old instances alive):
    // wait briefly, then skip if any instance already queued the retry notice in this session.
    try {
      await new Promise((resolve) => setTimeout(resolve, 250 + Math.floor(Math.random() * 250)))
      const fresh: any = await withTimeout(ctx.session.context({ sessionID }), 5000, "dedup-context")
      const messages2: any[] = Array.isArray(fresh) ? fresh : []
      const idx = messages2.findIndex((m: any) => m?.id === lastAssistant.id)
      const noticeQueued = messages2
        .slice(idx + 1)
        .some((m: any) => m?.type === "user" && String(m?.text ?? "").startsWith("System notice: your previous message contained a tool call"))
      if (noticeQueued) {
        log(`leak: notice already queued by another instance, skipping (pid ${process.pid} ${lastAssistant.id})`)
        return
      }
    } catch (error) {
      log(`leak: dedup re-check failed, proceeding: ${String(error)}`)
    }

    st.consecutive += 1
    if (st.consecutive > LEAK_MAX_CONSECUTIVE) {
      log(`leak: gave up after ${LEAK_MAX_CONSECUTIVE} consecutive retries (${sessionID} ${lastAssistant.id}), counter reset`)
      st.consecutive = 0
      leakStates.set(sessionID, st)
      return
    }
    log(`leak: tool call leaked as text (pid ${process.pid} ${sessionID} ${lastAssistant.id}), retry ${st.consecutive}/${LEAK_MAX_CONSECUTIVE}`)
    await ctx.session.prompt({ sessionID, text: LEAK_RETRY_TEXT })
  } catch (error) {
    log(`leak: error: ${String(error)}`)
  } finally {
    st.inFlight = false
  }
}

// Regression fixture for the missed live leak (1604 chars): a shell call with a heredoc body,
// several times longer than the 1200-char cap that used to reject it.
const LEAK_SELFTEST_LONG = `<tool_call>\n<function=bash>\n<parameter=command>\ncat > /tmp/opencode/cas_start.cpp <<'EOF'\n#include "pithread.h"\n${"// padding line so the fixture stays long\n".repeat(120)}EOF\ng++ -o /tmp/opencode/cas_start /tmp/opencode/cas_start.cpp\n</parameter>\n<parameter=timeout>\n300000\n</parameter>\n</function>\n</tool_call>`

function leakSelfTest() {
  const cases: Array<[string, string, boolean]> = [
    ["prefix-leak", "<tool_call>\n<function=bash>\n<parameter=command>\necho test\n</parameter>\n</function>\n</tool_call>", true],
    ["unknown-fn-leak", "<tool_call>\n<function=<unknown>\n</function>\n</tool_call>", true],
    ["preamble-suffix-leak", "I'll run that now.\n<tool_call>\n<function=bash>\n<parameter=command>\nls\n</parameter>\n</function>\n</tool_call>", true],
    ["normal-text", "Done. Build passed, tests are green.", false],
    ["empty", "   ", false],
    ["tag-in-middle", `We discussed the ${"<tool_call>"} tag earlier in this session.`, false],
    ["fenced-example", "Example format:\n```\n<tool_call>\n<function=bash>\n</function>\n</tool_call>\n```\nKeep it short.", false],
    ["fenced-example-suffix", "The call must be closed like this:\n```\n</tool_call>\n```", false],
    ["long-leak-5kB", LEAK_SELFTEST_LONG, LEAK_MAX_LEN === 0 || LEAK_SELFTEST_LONG.length <= LEAK_MAX_LEN],
    ["leak-with-fence-in-command", "<tool_call>\n<function=bash>\n<parameter=command>\nprintf '```md\nnote\n```' > /tmp/n.md\n</parameter>\n</function>\n</tool_call>", true],
    ["inline-code-suffix-untouched", "Wrap every call in a closing `</tool_call>` and nothing else.", false],
  ]
  if (LEAK_MAX_LEN > 0) cases.push([`cap-exceeded-${LEAK_MAX_LEN}`, `<tool_call>${"x".repeat(LEAK_MAX_LEN + 10)}`, false])
  for (const [name, text, expected] of cases) {
    const actual = leakedToolCallText(text)
    log(`leak self-test ${name}: ${actual === expected ? "PASS" : `FAIL got=${actual}`}`)
  }
}

// ============================================================
// plugin
// ============================================================

export default Plugin.define({
  id: "andrey.opencode-fixes",
  async setup(ctx) {
    // --- bool-args ---
    let editorGet: ((id: string) => any) | undefined
    await ctx.tool.transform((draft) => {
      editorGet = (id: string) => draft.get(id)
    })

    try {
      const tools = await ctx.tool.list()
      let tested = 0
      for (const tool of tools as any[]) {
        const resolved = resolveCached(tool.id, tool)
        if (tool.id === "shell" && resolved) {
          boolSelfTest(resolved)
          tested = 1
        }
        if (process.env.OPENCODE_FIXES_DEBUG && !resolved) {
          log(`bool: no schema resolved for tool ${tool.id}`)
        }
      }
      log(`bool: schema resolvers installed, shell self-test ${tested ? "executed" : "SKIPPED"}`)
    } catch (error) {
      log(`bool: setup probe failed: ${String(error)}`)
    }

    await ctx.tool.hook("execute.before", (event: any) => {
      try {
        const resolved = resolveCached(event.tool, editorGet?.(event.tool))
        if (!resolved) return
        const fixed = coerceInput(event.input, resolved)
        if (fixed !== event.input) {
          event.input = fixed
          log(`bool: repaired ${event.tool}: ${JSON.stringify(fixed)}`)
        } else if (process.env.OPENCODE_FIXES_DEBUG) {
          log(`bool: hook ${event.tool}: no changes`)
        }
      } catch (error) {
        log(`bool: hook ${event.tool} error: ${String(error)}`)
      }
    })

    // --- toolcall-leak ---
    leakSelfTest()
    const leakRetryLoop = async () => {
      while (true) {
        try {
          const stream: any = await ctx.event.subscribe()
          for await (const event of stream) {
            if (event?.type !== "session.execution.succeeded") continue
            const sessionID = event?.data?.sessionID
            if (typeof sessionID === "string") void handleTurnEnd(ctx, sessionID)
          }
          log("leak: event stream ended, resubscribing")
        } catch (error) {
          log(`leak: stream error: ${String(error)}`)
        }
        await new Promise((resolve) => setTimeout(resolve, 5000))
      }
    }
    void leakRetryLoop()
  },
})