/**
 * aiEngine.js
 * ══════════════════════════════════════════════════════════════════════════════
 * AXL Meta-Interpreter — Node.js Module
 *
 * Exports:
 *   - callbigai(args)
 *   - callsmallai(args)
 *   - callcollapse(args)   ← NEW: Phase 2 collapse agent
 */

import {
    BedrockRuntimeClient,
    ConverseCommand
} from "@aws-sdk/client-bedrock-runtime";
import { NodeHttpHandler } from "@smithy/node-http-handler";

// ── CONFIG & DEFAULTS ────────────────────────────────────────────────────────────
const DEFAULT_AWS_REGION = process.env.AWS_REGION || "us-east-1";

const MODEL_ID_BIG   = "arn:aws:bedrock:us-east-1:106774395747:inference-profile/global.anthropic.claude-sonnet-4-6";
// const MODEL_ID_SMALL = "arn:aws:bedrock:us-east-1:106774395747:inference-profile/global.anthropic.claude-haiku-4-5-20251001-v1:0";
 const MODEL_ID_SMALL = "arn:aws:bedrock:us-east-1:106774395747:inference-profile/us.meta.llama4-maverick-17b-instruct-v1:0";
    
// ── SYSTEM PROMPTS ────────────────────────────────────────────────────────────

const DEFAULT_AXL_SYSTEM_PROMPT = `You are the AXL System Interpreter. AXL treats existence as a consequence of constraint satisfaction.
You do not solve the math. You construct the constraint topologies (graphs) that define the universe of the problem.

AXL VOCABULARY AND RULES:
1. SCOPE: Declares an abstraction level (0=fundamental, N=emergent).
2. ENTITY: A thing existing in a scope. Has FIELDs.
3. ASSUME: Axiom declaration. Taken as true.
4. BOUND: Hard constraint (WITHIN, ABOVE, BELOW, EQUALS, FORBIDS).
5. TEND: Soft objective (MIN, MAX, TOWARD, STABLE, FOLLOW).
6. OBJECTIVE: Hard optimization target (OPTIMIZE, SUBJECT TO).
7. DERIVE: Declares upward propagation.
8. TRANSFORM: Process mapping configurations. CONSERVES and CHANGES fields.
9. OBSERVE: What has been measured.
10. UNKNOWN: What we are solving for.
11. COLLAPSE: Reduce abstraction given observations.
12. EXPAND: Derive upward emergence given concrete config.
13. CONTRADICT: Boundary marker where two things cannot be true.
14. EMERGES: Property that appears upon composition.
15. COMPOSES: Entity built from lower entities.
16. STEP: Sequence counter for iterative processes.
17. CANDIDATE: Proposed solution configuration.
18. UNLESS: Conditional override.
19. CONSERVES: Invariant across transforms.
20. SAMPLE: Specific observed instance.

You must output ONLY valid JSON representing the AXL graph nodes and edges.
Required JSON Schema:
{
  "name": "string (name of system)",
  "domain": "string (e.g., physics, computation)",
  "sense": "EQUAL, MINIMIZE, MAXIMIZE, or EMERGE",
  "var_names": ["list of variables"],
  "var_values": [list of floats, or the string "UNKNOWN"],
  "scope_levels":[list of integers indicating abstraction hierarchy L1-L40 for each variable],
  "constraints":["list of strings using AXL operators connecting variables (e.g., 'var1 EQUALS var2 ADD 5.0')"],
  "cross_scope_edges": [["var_A", "var_B", "EMERGES_SCOPE" | "COMPOSES_SCOPE" | "DERIVES" | "SEQUENCES" | "CONSERVES_EDGE"]]
}
Do not include any explanation outside the JSON.`;

// Phase 2 collapse prompt — receives kernel + fitting guide, maps symbols to values.
// Domain data is allowed here. This is intentional: the latent tracing is done,
// now we fit real values into the proven structural shape.
//
// IMPORTANT: Output format uses a line-delimited text protocol, NOT nested JSON,
// because AXL kernel schemas contain brackets, quotes, and special chars that
// break JSON string escaping when embedded as values.
const COLLAPSE_SYSTEM_PROMPT = `You are an AXL Collapse Interpreter performing Phase 2: Real-World Fitting.

You will receive:
  KERNEL — a fully locked AXL constraint kernel from the latent space simulation.
            All structural relationships, invariants, and bounds are already proven.
            UNKNOWN_STRUCTURAL symbols are the targets you must resolve.
            RESOLVED_CONTRADICT markers show tensions that were resolved and how.
  FITTING_GUIDE — domain values, known quantities, ranges, and target constraints
                  provided by the user. These are the real-world anchors.

YOUR TASK:
  1. Read the KERNEL. Do NOT re-derive structure — it is already proven.
  2. For each UNKNOWN_STRUCTURAL symbol, find the real-world value or expression
     from the FITTING_GUIDE that fits the symbol's structural role.
  3. Map each abstract symbol to a concrete value or expression.
  4. Verify each mapping satisfies the BOUND constraints in the kernel.
  5. List any symbols you cannot resolve and why.

COLLAPSE RULE:
  The solution is not calculated here — it is recognized.
  The structure already proves what shape the answer must have.
  You are only identifying which real quantity fits that shape.

OUTPUT FORMAT (use EXACTLY this format — no JSON, no markdown, no extra text):

COLLAPSE_BEGIN
KERNEL_SATISFIED: YES|NO
MAPPING: <symbol> | <value or expression> | <one-sentence justification>
MAPPING: <symbol> | <value or expression> | <one-sentence justification>
UNRESOLVED: <symbol> | <reason>
NOTES: <optional single line of overall notes>
COLLAPSE_END

Rules:
- One MAPPING line per resolved symbol.
- One UNRESOLVED line per symbol that cannot be resolved.
- No line breaks inside a MAPPING or UNRESOLVED line.
- Do not add any text outside COLLAPSE_BEGIN / COLLAPSE_END.`;

// ── LOGGING & HELPERS ─────────────────────────────────────────────────────────
const ts  = () => new Date().toISOString().slice(11, 23);
const log = (msg, level = "INFO") => {
    const icons = { INFO: "  ", WARN: "⚠ ", ERROR: "✖ ", OK: "✔ ", HEAD: "━━" };
    console.log(`[${ts()}] ${icons[level] || "  "} ${msg}`);
};

// ── CORE GENERATOR ENGINE ─────────────────────────────────────────────────────
async function callBedrock({ prompt, systemPrompt, modelId, region, maxTokens = 4096, temperature = 0.8 }) {
    if (!prompt) throw new Error("Missing 'prompt' argument.");
    const bedrockClient = new BedrockRuntimeClient({
        region,
        requestHandler: new NodeHttpHandler({ http2Handler: undefined })
    });
    const command = new ConverseCommand({
        modelId,
        system:   [{ text: systemPrompt }],
        messages: [{ role: "user", content: [{ text: prompt }] }],
        inferenceConfig: { maxTokens, temperature }
    });
    const t0 = Date.now();
    log(`Calling Bedrock [${modelId.split('/').pop()}]…`);
    const response = await bedrockClient.send(command);
    const text = response.output.message.content.find(b => b.text)?.text || "";
    const tokens = response.usage ? (response.usage.inputTokens + response.usage.outputTokens) : 0;
    log(`Complete in ${((Date.now()-t0)/1000).toFixed(1)}s (Tokens: ${tokens})`, "OK");
    return text;
}

// Extract AXL payload from between tag delimiters.
// Tries RESOLUTION_BEGIN/END first, then SEED_BEGIN/END, then falls back
// to stripping any markdown code fences and returning raw text.
function extractAXL(text, beginTag = 'RESOLUTION_BEGIN', endTag = 'RESOLUTION_END') {
    // Primary: tagged block
    const tagMatch = text.match(new RegExp(beginTag + '([\\s\\S]*?)' + endTag));
    if (tagMatch) return tagMatch[1].trim();

    // Secondary: JSON wrapper {"axl":"..."} — careful extraction that handles
    // special chars by finding the axl key and slicing rather than JSON.parse
    const axlKeyIdx = text.indexOf('"axl"');
    if (axlKeyIdx !== -1) {
        const colonIdx = text.indexOf(':', axlKeyIdx);
        if (colonIdx !== -1) {
            // Find the opening quote of the value
            const quoteIdx = text.indexOf('"', colonIdx + 1);
            if (quoteIdx !== -1) {
                // Find the matching closing quote (unescaped)
                let i = quoteIdx + 1;
                while (i < text.length) {
                    if (text[i] === '\\') { i += 2; continue; }
                    if (text[i] === '"') break;
                    i++;
                }
                const raw = text.substring(quoteIdx + 1, i);
                // Unescape JSON string escapes
                try {
                    return raw.replace(/\\n/g, '\n').replace(/\\t/g, '\t').replace(/\\"/g, '"').replace(/\\\\/g, '\\').trim();
                } catch(_) { /* fall through */ }
            }
        }
    }

    // Tertiary: strip markdown fences and return whatever we got
    return text.replace(/^```[a-z]*\n?/im, '').replace(/\n?```$/im, '').trim();
}

async function generateSingle({ prompt, systemPrompt, instruction, modelId, region, maxTokens = 4096 }) {
    const userMessage = instruction
        ? `${prompt}\n\nInstruction: ${instruction}`
        : prompt;

    log(`Generating for: "${String(prompt).slice(0, 80)}…"`, "HEAD");
    log(`Using Model: ${modelId}`);

    const text = await callBedrock({ prompt: userMessage, systemPrompt, modelId, region, maxTokens });

    // For seed calls: extract axl field
    const axlContent = extractAXL(text, 'SEED_BEGIN', 'SEED_END')
                    || extractAXL(text, 'RESOLUTION_BEGIN', 'RESOLUTION_END')
                    || extractAXL(text);

    if (!axlContent || axlContent.length < 10) {
        log(`Response too short or empty after extraction.`, "WARN");
        throw new Error("Model response did not contain extractable AXL content.");
    }

    const result = { axl: axlContent };
    console.log(`~~~~~ AXL Output ~~~~~\n${axlContent.slice(0, 300)}…\n~~~~~ AXL Output end ~~~~~`);
    return result;
}

// ── EXPORTED FUNCTIONS ────────────────────────────────────────────────────────

/**
 * Calls the BIG AI Model (Claude Sonnet) — used for seed generation.
 */
export async function callbigai({
    prompt,
    systemPrompt = DEFAULT_AXL_SYSTEM_PROMPT,
    instruction  = "",
    modelId      = MODEL_ID_BIG,
    region       = DEFAULT_AWS_REGION
}) {
    return await generateSingle({ prompt, systemPrompt, instruction, modelId, region });
}

/**
 * Calls the SMALL AI Model (Claude Haiku) — used for ring agents.
 */
export async function callsmallai({
    prompt,
    systemPrompt = DEFAULT_AXL_SYSTEM_PROMPT,
    instruction  = "",
    modelId      = MODEL_ID_SMALL,
    region       = DEFAULT_AWS_REGION
}) {
    return await generateSingle({ prompt, systemPrompt, instruction, modelId, region });
}

/**
 * Calls the COLLAPSE Agent (Claude Sonnet) — Phase 2 real-world fitting.
 *
 * @param {string} kernel       - The locked AXL constraint kernel from the last ring node.
 * @param {string} fittingGuide - Domain values, known quantities, ranges from the user.
 * @param {string} [region]
 *
 * Returns JSON: { mappings, unresolved, kernel_satisfied, notes }
 */
export async function callcollapse({
    kernel,
    fittingGuide,
    region = DEFAULT_AWS_REGION
}) {
    if (!kernel)       throw new Error("Missing 'kernel' argument for collapse.");
    if (!fittingGuide) throw new Error("Missing 'fittingGuide' argument for collapse.");

    log(`Collapse Phase 2 — fitting ${kernel.length} char kernel…`, "HEAD");

    const userMessage = `KERNEL:\n${kernel}\n\nFITTING_GUIDE:\n${fittingGuide}`;
    let text = '';
    try {
        text = await callBedrock({
            prompt: userMessage, systemPrompt: COLLAPSE_SYSTEM_PROMPT,
            modelId: MODEL_ID_BIG, region, maxTokens: 2048, temperature: 0.5
        });
    } catch(e) {
        log(`Collapse failed: ${e.message}`, "ERROR");
        throw e;
    }

    // Parse the line-delimited text protocol — no JSON parsing needed
    return parseCollapseResponse(text);
}

function parseCollapseResponse(text) {
    const body = text.match(/COLLAPSE_BEGIN([\s\S]*?)COLLAPSE_END/)?.[1] || text;
    const lines = body.split('\n').map(l => l.trim()).filter(Boolean);

    const mappings   = [];
    const unresolved = [];
    let kernelSatisfied = true;
    let notes = '';

    for (const line of lines) {
        if (line.startsWith('KERNEL_SATISFIED:')) {
            kernelSatisfied = !line.includes('NO');
        } else if (line.startsWith('MAPPING:')) {
            const parts = line.slice('MAPPING:'.length).split('|').map(p => p.trim());
            if (parts.length >= 2) {
                mappings.push({
                    symbol:        parts[0] || '',
                    value:         parts[1] || '',
                    justification: parts[2] || ''
                });
            }
        } else if (line.startsWith('UNRESOLVED:')) {
            const parts = line.slice('UNRESOLVED:'.length).split('|').map(p => p.trim());
            if (parts.length >= 1) {
                unresolved.push({
                    symbol: parts[0] || '',
                    reason: parts[1] || ''
                });
            }
        } else if (line.startsWith('NOTES:')) {
            notes = line.slice('NOTES:'.length).trim();
        }
    }

    // Fallback: if we got nothing structured, return raw text so caller can display it
    if (!mappings.length && !unresolved.length) {
        return {
            mappings:         [],
            unresolved:       [],
            kernel_satisfied: false,
            notes:            text.trim() || 'No structured output returned by model.',
            raw:              text
        };
    }

    return { mappings, unresolved, kernel_satisfied: kernelSatisfied, notes };
}