export type TabularExtractionJsonObject = {
    column_index?: unknown;
    summary?: unknown;
    flag?: unknown;
    reasoning?: unknown;
};

/** Escape literal newlines/tabs/CR inside JSON string literals so JSON.parse succeeds. */
export function escapeLiteralWhitespaceInJsonStrings(jsonStr: string): string {
    let cleaned = "";
    let inStr = false;
    let isEsc = false;
    for (let i = 0; i < jsonStr.length; i++) {
        const c = jsonStr[i];
        if (inStr) {
            if (isEsc) {
                cleaned += c;
                isEsc = false;
            } else if (c === "\\") {
                cleaned += c;
                isEsc = true;
            } else if (c === '"') {
                cleaned += c;
                inStr = false;
            } else if (c === "\n") {
                cleaned += "\\n";
            } else if (c === "\r") {
                cleaned += "\\r";
            } else if (c === "\t") {
                cleaned += "\\t";
            } else {
                cleaned += c;
            }
        } else {
            if (c === '"') {
                inStr = true;
            }
            cleaned += c;
        }
    }
    return cleaned;
}

function findCompleteJsonObjectEnd(text: string, startIndex: number): number {
    let depth = 0;
    let inString = false;
    let escaped = false;

    for (let i = startIndex; i < text.length; i++) {
        const char = text[i];
        if (inString) {
            if (escaped) escaped = false;
            else if (char === "\\") escaped = true;
            else if (char === '"') inString = false;
        } else {
            if (char === '"') inString = true;
            else if (char === "{") depth++;
            else if (char === "}") {
                depth--;
                if (depth === 0) return i;
            }
        }
    }
    return -1;
}

export function parseTabularExtractionObject(
    jsonStr: string,
): TabularExtractionJsonObject | null {
    try {
        const cleaned = escapeLiteralWhitespaceInJsonStrings(jsonStr);
        return JSON.parse(cleaned) as TabularExtractionJsonObject;
    } catch {
        return null;
    }
}

/** Extract complete top-level JSON objects from a text buffer. */
export function drainTabularJsonObjects(buffer: string): {
    objects: TabularExtractionJsonObject[];
    remainder: string;
    parseFailures: number;
} {
    const objects: TabularExtractionJsonObject[] = [];
    let remainder = buffer;
    let parseFailures = 0;

    let startIndex = remainder.indexOf("{");
    while (startIndex !== -1) {
        const endIndex = findCompleteJsonObjectEnd(remainder, startIndex);
        if (endIndex === -1) break;

        const jsonStr = remainder.slice(startIndex, endIndex + 1);
        remainder = remainder.slice(endIndex + 1);

        const parsed = parseTabularExtractionObject(jsonStr);
        if (parsed) {
            objects.push(parsed);
        } else {
            parseFailures += 1;
        }

        startIndex = remainder.indexOf("{");
    }

    return { objects, remainder, parseFailures };
}

/** Incrementally parse streamed model output into tabular extraction objects. */
export class TabularJsonStreamParser {
    private buffer = "";
    private totalParseFailures = 0;

    append(delta: string): TabularExtractionJsonObject[] {
        this.buffer += delta;
        return this.drain();
    }

    flush(): TabularExtractionJsonObject[] {
        return this.drain();
    }

    get remainder(): string {
        return this.buffer;
    }

    get parseFailureCount(): number {
        return this.totalParseFailures;
    }

    private drain(): TabularExtractionJsonObject[] {
        const { objects, remainder, parseFailures } = drainTabularJsonObjects(
            this.buffer,
        );
        this.buffer = remainder;
        this.totalParseFailures += parseFailures;
        return objects;
    }
}