// Scyla Dataset Module — Text Pair Training Data
//
// Load, create, edit, and save datasets of prompt→response pairs.
// Supports CSV, JSON lines, and simple pipe-delimited formats.
//
// Usage:
//   use dataset
//   let data = dataset_load("training_data.csv")
//   let count = dataset_len(data)
//   let pair = dataset_get(data, 0)   // [prompt, response]
//   dataset_save(data, "output.csv")
//
// Aislinn McCubbins, April 2026

// ── Core: A dataset is an Array of Arrays [[prompt, response], ...] ──

fn dataset_new() -> Array {
    return array_new()
}

fn dataset_add(data: Array, prompt: String, response: String) -> Array {
    let pair = array_new()
    pair = array_push(pair, prompt)
    pair = array_push(pair, response)
    data = array_push(data, pair)
    return data
}

fn dataset_len(data: Array) -> Int {
    return array_len(data)
}

fn dataset_get(data: Array, idx: Int) -> Array {
    return array_get(data, idx)
}

fn dataset_get_prompt(data: Array, idx: Int) -> String {
    let pair = array_get(data, idx)
    return array_get(pair, 0)
}

fn dataset_get_response(data: Array, idx: Int) -> String {
    let pair = array_get(data, idx)
    return array_get(pair, 1)
}

fn dataset_remove(data: Array, idx: Int) -> Array {
    // Remove row at index by building new array without it
    let result = array_new()
    let i = 0
    while i < array_len(data) {
        if i == idx {
            // skip
        } else {
            result = array_push(result, array_get(data, i))
        }
        i = i + 1
    }
    return result
}

fn dataset_set(data: Array, idx: Int, prompt: String, response: String) -> Array {
    let pair = array_new()
    pair = array_push(pair, prompt)
    pair = array_push(pair, response)
    return array_set(data, idx, pair)
}

// ── Loading ──

fn dataset_load(path: String) -> Array {
    let content = read_file(path)
    if starts_with(content, "error:") {
        println("Error loading dataset: " + content)
        return array_new()
    }

    // Detect format
    if str_ends_with(path, ".csv") {
        return dataset_parse_csv(content)
    }
    if str_ends_with(path, ".jsonl") {
        return dataset_parse_jsonl(content)
    }
    if str_ends_with(path, ".json") {
        return dataset_parse_json(content)
    }
    if str_ends_with(path, ".txt") {
        return dataset_parse_pipe(content)
    }
    // Default: try CSV
    return dataset_parse_csv(content)
}

fn dataset_parse_csv(content: String) -> Array {
    let data = array_new()
    let lines = str_lines(content)
    let i = 0
    while i < array_len(lines) {
        let line = array_get(lines, i)
        if i == 0 {
            // Skip header row if it looks like headers
            let lower = str_lower(line)
            if str_contains(lower, "prompt") {
                i = i + 1
            } else {
                // Not a header, parse it
                let pair = parse_csv_line(line)
                if array_len(pair) >= 2 {
                    data = dataset_add(data, array_get(pair, 0), array_get(pair, 1))
                }
                i = i + 1
            }
        } else {
            let pair = parse_csv_line(line)
            if array_len(pair) >= 2 {
                data = dataset_add(data, array_get(pair, 0), array_get(pair, 1))
            }
            i = i + 1
        }
    }
    return data
}

fn parse_csv_line(line: String) -> Array {
    // Handle quoted CSV fields: "prompt","response"
    let result = array_new()
    let trimmed = str_trim(line)
    if str_len(trimmed) == 0 { return result }

    if starts_with(trimmed, "\"") {
        // Quoted CSV
        let fields = array_new()
        let current = ""
        let in_quote = false
        let i = 0
        while i < str_len(trimmed) {
            let ch = str_char_at(trimmed, i)
            if str_eq(ch, "\"") {
                if in_quote {
                    // Check for escaped quote ""
                    if i + 1 < str_len(trimmed) {
                        let next = str_char_at(trimmed, i + 1)
                        if str_eq(next, "\"") {
                            current = current + "\""
                            i = i + 2
                        } else {
                            in_quote = false
                            i = i + 1
                        }
                    } else {
                        in_quote = false
                        i = i + 1
                    }
                } else {
                    in_quote = true
                    i = i + 1
                }
            } else {
                if str_eq(ch, ",") {
                    if in_quote {
                        current = current + ch
                    } else {
                        fields = array_push(fields, current)
                        current = ""
                    }
                    i = i + 1
                } else {
                    current = current + ch
                    i = i + 1
                }
            }
        }
        fields = array_push(fields, current)
        return fields
    } else {
        // Simple comma split
        return str_split(trimmed, ",")
    }
}

fn dataset_parse_pipe(content: String) -> Array {
    // Format: prompt|response (one per line)
    let data = array_new()
    let lines = str_lines(content)
    let i = 0
    while i < array_len(lines) {
        let line = str_trim(array_get(lines, i))
        if str_len(line) > 0 {
            let parts = str_split(line, "|")
            if array_len(parts) >= 2 {
                data = dataset_add(data, str_trim(array_get(parts, 0)), str_trim(array_get(parts, 1)))
            }
        }
        i = i + 1
    }
    return data
}

fn dataset_parse_jsonl(content: String) -> Array {
    // JSON Lines: {"prompt":"...", "response":"..."} per line
    let data = array_new()
    let lines = str_lines(content)
    let i = 0
    while i < array_len(lines) {
        let line = str_trim(array_get(lines, i))
        if str_len(line) > 0 {
            let prompt = extract_json_field(line, "prompt")
            let response = extract_json_field(line, "response")
            if str_len(prompt) == 0 {
                // Try "input" / "output" format
                prompt = extract_json_field(line, "input")
                response = extract_json_field(line, "output")
            }
            if str_len(prompt) == 0 {
                // Try "question" / "answer" format
                prompt = extract_json_field(line, "question")
                response = extract_json_field(line, "answer")
            }
            if str_len(prompt) > 0 {
                data = dataset_add(data, prompt, response)
            }
        }
        i = i + 1
    }
    return data
}

fn dataset_parse_json(content: String) -> Array {
    // JSON array: [{"prompt":"...", "response":"..."},...]
    // Treat it like JSONL by splitting on },{
    let cleaned = str_replace(content, "\},\{", "\}\n\{")
    return dataset_parse_jsonl(cleaned)
}

fn extract_json_field(json: String, field: String) -> String {
    let key = "\"" + field + "\""
    let pos = str_find(json, key)
    if pos < 0 { return "" }

    let after_key = str_slice(json, pos + str_len(key), str_len(json))
    let colon = str_find(after_key, ":")
    if colon < 0 { return "" }

    let after_colon = str_trim(str_slice(after_key, colon + 1, str_len(after_key)))
    if starts_with(after_colon, "\"") {
        let value_start = str_slice(after_colon, 1, str_len(after_colon))
        // Find closing unescaped quote
        let result = ""
        let vi = 0
        while vi < str_len(value_start) {
            let ch = str_char_at(value_start, vi)
            if str_eq(ch, "\\") {
                if vi + 1 < str_len(value_start) {
                    let next = str_char_at(value_start, vi + 1)
                    if str_eq(next, "n") {
                        result = result + "\n"
                    } else {
                        if str_eq(next, "\"") {
                            result = result + "\""
                        } else {
                            result = result + next
                        }
                    }
                    vi = vi + 2
                } else {
                    vi = vi + 1
                }
            } else {
                if str_eq(ch, "\"") {
                    return result
                } else {
                    result = result + ch
                    vi = vi + 1
                }
            }
        }
        return result
    }
    return ""
}

// ── Saving ──

fn dataset_save_csv(data: Array, path: String) -> String {
    let content = "prompt,response\n"
    let i = 0
    while i < array_len(data) {
        let pair = array_get(data, i)
        let prompt = str_replace(array_get(pair, 0), "\"", "\"\"")
        let response = str_replace(array_get(pair, 1), "\"", "\"\"")
        content = content + "\"" + prompt + "\",\"" + response + "\"\n"
        i = i + 1
    }
    return write_file(path, content)
}

fn dataset_save_jsonl(data: Array, path: String) -> String {
    let content = ""
    let i = 0
    while i < array_len(data) {
        let pair = array_get(data, i)
        let prompt = str_replace(str_replace(array_get(pair, 0), "\\", "\\\\"), "\"", "\\\"")
        let response = str_replace(str_replace(array_get(pair, 1), "\\", "\\\\"), "\"", "\\\"")
        content = content + "\{\"prompt\":\"" + prompt + "\",\"response\":\"" + response + "\"\}\n"
        i = i + 1
    }
    return write_file(path, content)
}

fn dataset_save_pipe(data: Array, path: String) -> String {
    let content = ""
    let i = 0
    while i < array_len(data) {
        let pair = array_get(data, i)
        content = content + array_get(pair, 0) + "|" + array_get(pair, 1) + "\n"
        i = i + 1
    }
    return write_file(path, content)
}

fn dataset_save(data: Array, path: String) -> String {
    if str_ends_with(path, ".csv") { return dataset_save_csv(data, path) }
    if str_ends_with(path, ".jsonl") { return dataset_save_jsonl(data, path) }
    if str_ends_with(path, ".txt") { return dataset_save_pipe(data, path) }
    // Default CSV
    return dataset_save_csv(data, path)
}

// ── Display ──

fn dataset_preview(data: Array, max_rows: Int) -> String {
    let result = "Dataset: " + int_to_str(array_len(data)) + " pairs\n"
    result = result + "─────────────────────────────────────────\n"
    let i = 0
    let show = max_rows
    if array_len(data) < show { show = array_len(data) }
    while i < show {
        let pair = array_get(data, i)
        let prompt = array_get(pair, 0)
        let response = array_get(pair, 1)
        // Truncate long values
        if str_len(prompt) > 40 {
            prompt = str_slice(prompt, 0, 40) + "..."
        }
        if str_len(response) > 40 {
            response = str_slice(response, 0, 40) + "..."
        }
        result = result + int_to_str(i + 1) + ". [P] " + prompt + "\n"
        result = result + "   [R] " + response + "\n"
        i = i + 1
    }
    if array_len(data) > show {
        result = result + "... and " + int_to_str(array_len(data) - show) + " more\n"
    }
    return result
}

// ── Merge / Split ──

fn dataset_merge(a: Array, b: Array) -> Array {
    let result = a
    let i = 0
    while i < array_len(b) {
        result = array_push(result, array_get(b, i))
        i = i + 1
    }
    return result
}

fn dataset_shuffle(data: Array) -> Array {
    // Simple Fisher-Yates using timestamp as seed
    let t = timestamp()
    let result = data
    let i = array_len(result) - 1
    while i > 0 {
        // Simple pseudo-random: use modulo of timestamp * index
        let j = (t * (i + 7)) % (i + 1)
        if j < 0 { j = 0 - j }
        if j > i { j = i }
        // Swap
        let tmp = array_get(result, i)
        result = array_set(result, i, array_get(result, j))
        result = array_set(result, j, tmp)
        i = i - 1
    }
    return result
}
