#!/usr/bin/env python3
"""Patch v2: Fix MemoryExtractor to use Ollama native API with think:false.

The OpenAI-compatible endpoint doesn't support think:false, causing Qwen3 14B
to spend all tokens on reasoning (content="", reasoning="..."). This patch
replaces the OpenAI SDK call with a direct fetch() to Ollama's native /api/chat
endpoint with think:false, which produces actual JSON output.
"""

import sys

PLUGIN_PATH = '/Users/aivagenesis/.npm-global/lib/node_modules/openclaw/extensions/memory-lancedb/index.ts'

with open(PLUGIN_PATH, 'r') as f:
    content = f.read()

# === PATCH: Replace MemoryExtractor class with native Ollama fetch version ===

OLD_CLASS_START = '''class MemoryExtractor {
  private client: OpenAI;
  private model: string;

  constructor(apiKey: string, baseUrl?: string, model?: string) {
    this.client = new OpenAI({ apiKey, ...(baseUrl ? { baseURL: baseUrl } : {}) });
    this.model = model ?? "qwen3:14b";
  }

  async extractFacts(
    conversation: string,
  ): Promise<Array<{ text: string; category: MemoryCategory; importance: number }>> {
    try {
      const systemPrompt = [
        "You are a memory extraction system. Extract facts worth remembering from this conversation.",
        "Output ONLY a JSON array of objects with these fields:",
        '- "text": standalone fact sentence (understandable without context)',
        '- "category": one of "preference", "fact", "decision", "entity", "other"',
        '- "importance": number 0.0 to 1.0',
        "",
        "Rules:",
        "- Extract 0 to 5 facts maximum",
        "- Only extract NEW, specific information (not common knowledge)",
        "- Names, dates, preferences, decisions, project details = high importance",
        "- Greetings, small talk, vague statements = DO NOT extract",
        "- Each fact MUST be self-contained (understandable alone)",
        "- Output ONLY valid JSON array, nothing else",
        "- If nothing is worth remembering, output []",
      ].join("\\n");

      const response = await this.client.chat.completions.create({
        model: this.model,
        messages: [
          { role: "system", content: systemPrompt },
          { role: "user", content: conversation.slice(0, 3000) },
        ],
        temperature: 0.1,
        max_tokens: 500,
      });

      const raw = response.choices?.[0]?.message?.content?.trim();
      if (!raw) return [];

      // Extract JSON array from response (handle markdown code blocks)
      const jsonMatch = raw.match(/\\[\\s*\\{[\\s\\S]*\\}\\s*\\]/);
      const jsonStr = jsonMatch ? jsonMatch[0] : raw;

      let parsed: unknown;
      try {
        parsed = JSON.parse(jsonStr);
      } catch {
        return [];
      }
      if (!Array.isArray(parsed)) return [];

      return (parsed as Array<Record<string, unknown>>)
        .filter(
          (item) =>
            item &&
            typeof item.text === "string" &&
            (item.text as string).length >= 10 &&
            (item.text as string).length <= 500,
        )
        .map((item) => ({
          text: String(item.text),
          category: MEMORY_CATEGORIES.includes(item.category as MemoryCategory)
            ? (item.category as MemoryCategory)
            : "fact",
          importance:
            typeof item.importance === "number"
              ? Math.min(1, Math.max(0, item.importance))
              : 0.7,
        }));
    } catch {
      return [];
    }
  }
}'''

NEW_CLASS = '''class MemoryExtractor {
  private ollamaUrl: string;
  private model: string;

  constructor(_apiKey: string, _baseUrl?: string, model?: string) {
    // Use Ollama native API (not OpenAI-compatible) to support think:false
    this.ollamaUrl = "http://localhost:11434/api/chat";
    this.model = model ?? "qwen3:14b";
  }

  async extractFacts(
    conversation: string,
  ): Promise<Array<{ text: string; category: MemoryCategory; importance: number }>> {
    try {
      const systemPrompt = [
        "You are a memory extraction system. Extract facts worth remembering from this conversation.",
        "Output ONLY a JSON array of objects with these fields:",
        '- "text": standalone fact sentence (understandable without context)',
        '- "category": one of "preference", "fact", "decision", "entity", "other"',
        '- "importance": number 0.0 to 1.0 (names/dates/decisions = 0.9, general facts = 0.7)',
        "",
        "Rules:",
        "- Extract 0 to 5 facts maximum",
        "- Only extract NEW, specific information (not common knowledge)",
        "- Names, dates, preferences, decisions, project details = high importance",
        "- Greetings, small talk, vague statements = DO NOT extract",
        "- Each fact MUST be self-contained (understandable alone)",
        "- Output ONLY valid JSON array, nothing else",
        "- If nothing is worth remembering, output []",
      ].join("\\n");

      const body = JSON.stringify({
        model: this.model,
        messages: [
          { role: "system", content: systemPrompt },
          { role: "user", content: conversation.slice(0, 3000) },
        ],
        stream: false,
        think: false,
        options: { temperature: 0.1, num_predict: 600 },
      });

      const resp = await fetch(this.ollamaUrl, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body,
        signal: AbortSignal.timeout(30000),
      });

      if (!resp.ok) return [];
      const data = (await resp.json()) as Record<string, unknown>;
      const msg = data.message as Record<string, unknown> | undefined;
      const raw = (typeof msg?.content === "string" ? msg.content : "").trim();
      if (!raw) return [];

      // Extract JSON array from response (handle markdown code blocks)
      const jsonMatch = raw.match(/\\[\\s*\\{[\\s\\S]*\\}\\s*\\]/);
      const jsonStr = jsonMatch ? jsonMatch[0] : raw;

      let parsed: unknown;
      try {
        parsed = JSON.parse(jsonStr);
      } catch {
        return [];
      }
      if (!Array.isArray(parsed)) return [];

      return (parsed as Array<Record<string, unknown>>)
        .filter(
          (item) =>
            item &&
            typeof item.text === "string" &&
            (item.text as string).length >= 10 &&
            (item.text as string).length <= 500,
        )
        .map((item) => ({
          text: String(item.text),
          category: MEMORY_CATEGORIES.includes(item.category as MemoryCategory)
            ? (item.category as MemoryCategory)
            : "fact",
          importance:
            typeof item.importance === "number"
              ? Math.min(1, Math.max(0, item.importance))
              : 0.7,
        }));
    } catch {
      return [];
    }
  }
}'''

if OLD_CLASS_START not in content:
    print('ERROR: Could not find existing MemoryExtractor class (already patched or missing)')
    print('Checking for v2 marker...')
    if 'ollamaUrl' in content:
        print('ALREADY PATCHED: v2 (native Ollama) class found')
        sys.exit(0)
    sys.exit(1)

content = content.replace(OLD_CLASS_START, NEW_CLASS, 1)
print('PATCH: Replaced MemoryExtractor with native Ollama fetch (think:false)')

# Write patched file
with open(PLUGIN_PATH, 'w') as f:
    f.write(content)

print(f'SUCCESS: Patch v2 applied to {PLUGIN_PATH}')
print(f'  File size: {len(content)} chars, {content.count(chr(10))} lines')