const std = @import("std");

/// A single change from the model output.
pub const Change = struct {
    original: []const u8,
    replacement: []const u8,
    marker: []const u8,
    reason: []const u8,
};

/// Parsed model output: clean text + structured changes.
pub const ParsedOutput = struct {
    clean_text: []const u8,
    changes: [MAX_CHANGES]Change,
    change_count: u32,

    pub const MAX_CHANGES = 32;

    pub fn changeSlice(self: *const ParsedOutput) []const Change {
        return self.changes[0..self.change_count];
    }
};

/// Parse model output that contains rewritten text followed by a ```json block.
/// Format:
///   <rewritten text>
///
///   ```json
///   {"changes":[...]}
///   ```
pub fn parseModelOutput(output: []const u8) ParsedOutput {
    var result = ParsedOutput{
        .clean_text = output,
        .changes = undefined,
        .change_count = 0,
    };

    // Initialize changes
    for (&result.changes) |*c| {
        c.* = .{ .original = "", .replacement = "", .marker = "", .reason = "" };
    }

    // Find the ```json fence
    const fence = "```json";
    const fence_pos = std.mem.indexOf(u8, output, fence);
    if (fence_pos) |pos| {
        // Clean text is everything before the fence, trimmed
        result.clean_text = std.mem.trimRight(u8, output[0..pos], " \n\r\t");

        // JSON is between ```json and ```
        const json_start = pos + fence.len;
        const json_end_fence = std.mem.indexOfPos(u8, output, json_start, "```");
        const json_text = if (json_end_fence) |end|
            std.mem.trim(u8, output[json_start..end], " \n\r\t")
        else
            std.mem.trim(u8, output[json_start..], " \n\r\t");

        // Parse the changes array with a simple scanner
        parseChangesJson(json_text, &result);
    }

    return result;
}

/// Simple JSON scanner for the changes array.
/// Looks for "original", "replacement", "marker", "reason" string values.
fn parseChangesJson(json: []const u8, result: *ParsedOutput) void {
    // Find each object in the changes array
    var i: usize = 0;
    while (i < json.len and result.change_count < ParsedOutput.MAX_CHANGES) {
        // Find next object start
        const obj_start = std.mem.indexOfPos(u8, json, i, "{");
        if (obj_start == null) break;

        const obj_end = std.mem.indexOfPos(u8, json, obj_start.?, "}");
        if (obj_end == null) break;

        const obj = json[obj_start.? .. obj_end.? + 1];

        var change = &result.changes[result.change_count];
        change.original = extractStringField(obj, "original");
        change.replacement = extractStringField(obj, "replacement");
        change.marker = extractStringField(obj, "marker");
        change.reason = extractStringField(obj, "reason");

        if (change.original.len > 0 or change.replacement.len > 0) {
            result.change_count += 1;
        }

        i = obj_end.? + 1;
    }
}

/// Extract a string value for a given key from a JSON object substring.
/// Returns the value without quotes. This is a simple scanner, not a full parser.
fn extractStringField(obj: []const u8, field: []const u8) []const u8 {
    // Look for "field":"value" or "field": "value"
    var search_buf: [64]u8 = undefined;
    const search_key = std.fmt.bufPrint(&search_buf, "\"{s}\"", .{field}) catch return "";

    const key_pos = std.mem.indexOf(u8, obj, search_key);
    if (key_pos == null) return "";

    // Find the colon after the key
    var pos = key_pos.? + search_key.len;
    while (pos < obj.len and (obj[pos] == ':' or obj[pos] == ' ')) : (pos += 1) {}

    if (pos >= obj.len or obj[pos] != '"') return "";
    pos += 1; // skip opening quote

    // Find closing quote (handle escaped quotes)
    const start = pos;
    while (pos < obj.len) {
        if (obj[pos] == '\\' and pos + 1 < obj.len) {
            pos += 2; // skip escape sequence
            continue;
        }
        if (obj[pos] == '"') break;
        pos += 1;
    }

    return obj[start..pos];
}

test "parse model output" {
    const output =
        \\The company has serious problems.
        \\
        \\```json
        \\{"changes":[{"original":"I've","replacement":"[removed]","marker":"contraction_high","reason":"Contraction marker"}]}
        \\```
    ;

    const parsed = parseModelOutput(output);
    try std.testing.expectEqualStrings("The company has serious problems.", parsed.clean_text);
    try std.testing.expect(parsed.change_count == 1);
    try std.testing.expectEqualStrings("I've", parsed.changeSlice()[0].original);
    try std.testing.expectEqualStrings("contraction_high", parsed.changeSlice()[0].marker);
}

test "parse output without json block" {
    const output = "Just some plain text without changes.";
    const parsed = parseModelOutput(output);
    try std.testing.expectEqualStrings(output, parsed.clean_text);
    try std.testing.expect(parsed.change_count == 0);
}