#!/usr/bin/env python3
"""Patch OpenClaw memory-lancedb plugin for LLM-based auto-capture.

Replaces regex-based shouldCapture with Qwen3 14B LLM extraction.
Keeps regex as fallback if LLM returns no results.
"""

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 1: Add MemoryExtractor class before Rule-based capture filter ===
marker1 = '// ============================================================================\n// Rule-based capture filter'
if marker1 not in content:
    print('ERROR: Could not find Rule-based capture filter section')
    sys.exit(1)

EXTRACTOR_CLASS = '''// ============================================================================
// LLM-based memory extraction
// ============================================================================

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 [];
    }
  }
}

'''

content = content.replace(marker1, EXTRACTOR_CLASS + marker1, 1)
print('PATCH 1: Added MemoryExtractor class')

# === PATCH 2: Add extractor initialization after embeddings ===
old_embed_init = 'const embeddings = new Embeddings(cfg.embedding.apiKey, cfg.embedding.model!, cfg.embedding.baseUrl);'
new_embed_init = old_embed_init + '\n    const extractor = new MemoryExtractor(cfg.embedding.apiKey, cfg.embedding.baseUrl, "qwen3:14b");'

if old_embed_init not in content:
    print('ERROR: Could not find embeddings initialization')
    sys.exit(1)

content = content.replace(old_embed_init, new_embed_init, 1)
print('PATCH 2: Added extractor initialization')

# === PATCH 3: Replace auto-capture handler ===
# Find the old handler by unique markers
old_start = '    // Auto-capture: analyze and store important information after agent ends'
old_end_pattern = "          api.logger.warn(`memory-lancedb: capture failed: ${String(err)}`);\n        }\n      });\n    }"

if old_start not in content:
    # Check if already patched
    if 'LLM-based extraction' in content:
        print('PATCH 3: Already patched (LLM-based extraction found)')
    else:
        print('ERROR: Could not find auto-capture start marker')
        sys.exit(1)
else:
    start_idx = content.index(old_start)
    end_idx = content.index(old_end_pattern, start_idx)
    end_idx += len(old_end_pattern)

    NEW_CAPTURE = '''    // Auto-capture: LLM-based extraction of memorable facts after agent ends
    if (cfg.autoCapture) {
      api.on("agent_end", async (event) => {
        if (!event.success || !event.messages || event.messages.length === 0) {
          return;
        }

        try {
          // Build conversation transcript from ALL messages (user + assistant)
          const parts: string[] = [];
          for (const msg of event.messages) {
            if (!msg || typeof msg !== "object") continue;
            const msgObj = msg as Record<string, unknown>;
            const role = String(msgObj.role ?? "unknown");
            const msgContent = msgObj.content;
            let text = "";
            if (typeof msgContent === "string") {
              text = msgContent;
            } else if (Array.isArray(msgContent)) {
              text = (msgContent as Array<Record<string, unknown>>)
                .filter((b) => b?.type === "text" && typeof b.text === "string")
                .map((b) => b.text as string)
                .join(" ");
            }
            if (text && text.length > 5) {
              parts.push(`[${role}]: ${text}`);
            }
          }

          const conversation = parts.join("\\n").slice(0, 4000);
          if (conversation.length < 20) return;

          // Primary path: LLM-based fact extraction via local Qwen3 14B
          const facts = await extractor.extractFacts(conversation);

          if (facts.length === 0) {
            // Fallback: regex-based capture for user messages only
            const userTexts: string[] = [];
            for (const msg of event.messages) {
              if (!msg || typeof msg !== "object") continue;
              const msgObj = msg as Record<string, unknown>;
              if (msgObj.role !== "user") continue;
              const c = msgObj.content;
              if (typeof c === "string") {
                userTexts.push(c);
              } else if (Array.isArray(c)) {
                for (const b of c as Array<Record<string, unknown>>) {
                  if (b?.type === "text" && typeof b.text === "string") {
                    userTexts.push(b.text as string);
                  }
                }
              }
            }
            const toCapture = userTexts.filter(
              (text) => text && shouldCapture(text, { maxChars: cfg.captureMaxChars }),
            );
            for (const text of toCapture.slice(0, 3)) {
              const category = detectCategory(text);
              const vector = await embeddings.embed(text);
              const existing = await db.search(vector, 1, 0.95);
              if (existing.length > 0) continue;
              await db.store({ text, vector, importance: 0.7, category });
            }
            return;
          }

          // Store LLM-extracted facts (limit 5 per conversation)
          let stored = 0;
          for (const fact of facts.slice(0, 5)) {
            const vector = await embeddings.embed(fact.text);
            // Dedup check (slightly lower threshold — LLM may rephrase)
            const existing = await db.search(vector, 1, 0.90);
            if (existing.length > 0) continue;
            await db.store({
              text: fact.text,
              vector,
              importance: fact.importance,
              category: fact.category,
            });
            stored++;
          }

          if (stored > 0) {
            api.logger.info(`memory-lancedb: auto-captured ${stored} memories (LLM-extracted)`);
          }
        } catch (err) {
          api.logger.warn(`memory-lancedb: capture failed: ${String(err)}`);
        }
      });
    }'''

    content = content[:start_idx] + NEW_CAPTURE + content[end_idx:]
    print('PATCH 3: Replaced auto-capture with LLM-based extraction')

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

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