// JSON Serialization and Deserialization for Lux
//
// Provides type-safe JSON encoding and decoding with support
// for schema versioning and custom codecs.
//
// Usage:
//   let json = Json.encode(user)  // Serialize to JSON string
//   let user = Json.decode(json)  // Deserialize from JSON string

// ============================================================
// JSON Value Type
// ============================================================

// Represents any JSON value
type JsonValue =
    | JsonNull
    | JsonBool(Bool)
    | JsonInt(Int)
    | JsonFloat(Float)
    | JsonString(String)
    | JsonArray(List<JsonValue>)
    | JsonObject(List<(String, JsonValue)>)

// ============================================================
// Encoding Primitives
// ============================================================

// Escape a string for JSON
pub fn escapeString(s: String): String = {
    let escaped = String.replace(s, "\\", "\\\\")
    let escaped = String.replace(escaped, "\"", "\\\"")
    let escaped = String.replace(escaped, "\n", "\\n")
    let escaped = String.replace(escaped, "\r", "\\r")
    let escaped = String.replace(escaped, "\t", "\\t")
    escaped
}

// Encode a JsonValue to a JSON string
pub fn encode(value: JsonValue): String = {
    match value {
        JsonNull => "null",
        JsonBool(b) => if b then "true" else "false",
        JsonInt(n) => toString(n),
        JsonFloat(f) => toString(f),
        JsonString(s) => "\"" + escapeString(s) + "\"",
        JsonArray(items) => {
            let encodedItems = List.map(items, encode)
            "[" + String.join(encodedItems, ",") + "]"
        },
        JsonObject(fields) => {
            let encodedFields = List.map(fields, fn(field: (String, JsonValue)): String => {
                match field {
                    (key, val) => "\"" + escapeString(key) + "\":" + encode(val)
                }
            })
            "{" + String.join(encodedFields, ",") + "}"
        }
    }
}

// Pretty-print a JsonValue with indentation
pub fn encodePretty(value: JsonValue): String =
    encodePrettyIndent(value, 0)

fn encodePrettyIndent(value: JsonValue, indent: Int): String = {
    let spaces = String.repeat("  ", indent)
    let nextSpaces = String.repeat("  ", indent + 1)

    match value {
        JsonNull => "null",
        JsonBool(b) => if b then "true" else "false",
        JsonInt(n) => toString(n),
        JsonFloat(f) => toString(f),
        JsonString(s) => "\"" + escapeString(s) + "\"",
        JsonArray(items) => {
            if List.length(items) == 0 then "[]"
            else {
                let encodedItems = List.map(items, fn(item: JsonValue): String =>
                    nextSpaces + encodePrettyIndent(item, indent + 1)
                )
                "[\n" + String.join(encodedItems, ",\n") + "\n" + spaces + "]"
            }
        },
        JsonObject(fields) => {
            if List.length(fields) == 0 then "{}"
            else {
                let encodedFields = List.map(fields, fn(field: (String, JsonValue)): String => {
                    match field {
                        (key, val) => nextSpaces + "\"" + escapeString(key) + "\": " + encodePrettyIndent(val, indent + 1)
                    }
                })
                "{\n" + String.join(encodedFields, ",\n") + "\n" + spaces + "}"
            }
        }
    }
}

// ============================================================
// Type-specific Encoders
// ============================================================

// Encode primitives
pub fn encodeNull(): JsonValue = JsonNull
pub fn encodeBool(b: Bool): JsonValue = JsonBool(b)
pub fn encodeInt(n: Int): JsonValue = JsonInt(n)
pub fn encodeFloat(f: Float): JsonValue = JsonFloat(f)
pub fn encodeString(s: String): JsonValue = JsonString(s)

// Encode a list
pub fn encodeList<A>(items: List<A>, encodeItem: fn(A): JsonValue): JsonValue =
    JsonArray(List.map(items, encodeItem))

// Encode an optional value
pub fn encodeOption<A>(opt: Option<A>, encodeItem: fn(A): JsonValue): JsonValue =
    match opt {
        None => JsonNull,
        Some(value) => encodeItem(value)
    }

// Encode a Result
pub fn encodeResult<T, E>(
    result: Result<T, E>,
    encodeOk: fn(T): JsonValue,
    encodeErr: fn(E): JsonValue
): JsonValue =
    match result {
        Ok(value) => JsonObject([
            ("ok", encodeOk(value))
        ]),
        Err(error) => JsonObject([
            ("error", encodeErr(error))
        ])
    }

// ============================================================
// Object Building Helpers
// ============================================================

// Create an empty JSON object
pub fn object(): JsonValue = JsonObject([])

// Add a field to a JSON object
pub fn withField(obj: JsonValue, key: String, value: JsonValue): JsonValue =
    match obj {
        JsonObject(fields) => JsonObject(List.concat(fields, [(key, value)])),
        _ => obj  // Not an object, return unchanged
    }

// Add a string field
pub fn withString(obj: JsonValue, key: String, value: String): JsonValue =
    withField(obj, key, JsonString(value))

// Add an int field
pub fn withInt(obj: JsonValue, key: String, value: Int): JsonValue =
    withField(obj, key, JsonInt(value))

// Add a bool field
pub fn withBool(obj: JsonValue, key: String, value: Bool): JsonValue =
    withField(obj, key, JsonBool(value))

// Add an optional field (only adds if Some)
pub fn withOptional<A>(obj: JsonValue, key: String, opt: Option<A>, encodeItem: fn(A): JsonValue): JsonValue =
    match opt {
        None => obj,
        Some(value) => withField(obj, key, encodeItem(value))
    }

// ============================================================
// Decoding
// ============================================================

// Result type for parsing
type ParseResult<A> = Result<A, String>

// Get a field from a JSON object
pub fn getField(obj: JsonValue, key: String): Option<JsonValue> =
    match obj {
        JsonObject(fields) => findField(fields, key),
        _ => None
    }

fn findField(fields: List<(String, JsonValue)>, key: String): Option<JsonValue> =
    match List.head(fields) {
        None => None,
        Some(field) => match field {
            (k, v) => if k == key then Some(v)
                     else findField(Option.getOrElse(List.tail(fields), []), key)
        }
    }

// Get a string field
pub fn getString(obj: JsonValue, key: String): Option<String> =
    match getField(obj, key) {
        Some(JsonString(s)) => Some(s),
        _ => None
    }

// Get an int field
pub fn getInt(obj: JsonValue, key: String): Option<Int> =
    match getField(obj, key) {
        Some(JsonInt(n)) => Some(n),
        _ => None
    }

// Get a bool field
pub fn getBool(obj: JsonValue, key: String): Option<Bool> =
    match getField(obj, key) {
        Some(JsonBool(b)) => Some(b),
        _ => None
    }

// Get an array field
pub fn getArray(obj: JsonValue, key: String): Option<List<JsonValue>> =
    match getField(obj, key) {
        Some(JsonArray(items)) => Some(items),
        _ => None
    }

// Get an object field
pub fn getObject(obj: JsonValue, key: String): Option<JsonValue> =
    match getField(obj, key) {
        Some(JsonObject(_) as obj) => Some(obj),
        _ => None
    }

// ============================================================
// Simple JSON Parser
// ============================================================

// Parse a JSON string into a JsonValue
// Note: This is a simplified parser for common cases
pub fn parse(json: String): Result<JsonValue, String> =
    parseValue(String.trim(json), 0).mapResult(fn(r: (JsonValue, Int)): JsonValue => {
        match r { (value, _) => value }
    })

fn parseValue(json: String, pos: Int): Result<(JsonValue, Int), String> = {
    let c = String.charAt(json, pos)
    match c {
        "n" => parseNull(json, pos),
        "t" => parseTrue(json, pos),
        "f" => parseFalse(json, pos),
        "\"" => parseString(json, pos),
        "[" => parseArray(json, pos),
        "{" => parseObject(json, pos),
        "-" => parseNumber(json, pos),
        _ => if isDigit(c) then parseNumber(json, pos)
             else if c == " " || c == "\n" || c == "\r" || c == "\t" then
                 parseValue(json, pos + 1)
             else Err("Unexpected character at position " + toString(pos))
    }
}

fn parseNull(json: String, pos: Int): Result<(JsonValue, Int), String> =
    if String.substring(json, pos, pos + 4) == "null" then Ok((JsonNull, pos + 4))
    else Err("Expected 'null' at position " + toString(pos))

fn parseTrue(json: String, pos: Int): Result<(JsonValue, Int), String> =
    if String.substring(json, pos, pos + 4) == "true" then Ok((JsonBool(true), pos + 4))
    else Err("Expected 'true' at position " + toString(pos))

fn parseFalse(json: String, pos: Int): Result<(JsonValue, Int), String> =
    if String.substring(json, pos, pos + 5) == "false" then Ok((JsonBool(false), pos + 5))
    else Err("Expected 'false' at position " + toString(pos))

fn parseString(json: String, pos: Int): Result<(JsonValue, Int), String> = {
    // Skip opening quote
    let start = pos + 1
    let result = parseStringContent(json, start, "")
    result.mapResult(fn(r: (String, Int)): (JsonValue, Int) => {
        match r { (s, endPos) => (JsonString(s), endPos) }
    })
}

fn parseStringContent(json: String, pos: Int, acc: String): Result<(String, Int), String> = {
    let c = String.charAt(json, pos)
    if c == "\"" then Ok((acc, pos + 1))
    else if c == "\\" then {
        let nextC = String.charAt(json, pos + 1)
        let escaped = match nextC {
            "n" => "\n",
            "r" => "\r",
            "t" => "\t",
            "\"" => "\"",
            "\\" => "\\",
            _ => nextC
        }
        parseStringContent(json, pos + 2, acc + escaped)
    }
    else if c == "" then Err("Unterminated string")
    else parseStringContent(json, pos + 1, acc + c)
}

fn parseNumber(json: String, pos: Int): Result<(JsonValue, Int), String> = {
    let result = parseNumberDigits(json, pos, "")
    result.mapResult(fn(r: (String, Int)): (JsonValue, Int) => {
        match r {
            (numStr, endPos) => {
                // Check if it's a float
                if String.contains(numStr, ".") then
                    (JsonFloat(parseFloat(numStr)), endPos)
                else
                    (JsonInt(parseInt(numStr)), endPos)
            }
        }
    })
}

fn parseNumberDigits(json: String, pos: Int, acc: String): Result<(String, Int), String> = {
    let c = String.charAt(json, pos)
    if isDigit(c) || c == "." || c == "-" || c == "e" || c == "E" || c == "+" then
        parseNumberDigits(json, pos + 1, acc + c)
    else if acc == "" then Err("Expected number at position " + toString(pos))
    else Ok((acc, pos))
}

fn parseArray(json: String, pos: Int): Result<(JsonValue, Int), String> = {
    // Skip opening bracket and whitespace
    let startPos = skipWhitespace(json, pos + 1)

    if String.charAt(json, startPos) == "]" then Ok((JsonArray([]), startPos + 1))
    else parseArrayItems(json, startPos, [])
}

fn parseArrayItems(json: String, pos: Int, acc: List<JsonValue>): Result<(JsonValue, Int), String> = {
    match parseValue(json, pos) {
        Err(e) => Err(e),
        Ok((value, nextPos)) => {
            let newAcc = List.concat(acc, [value])
            let afterWhitespace = skipWhitespace(json, nextPos)
            let c = String.charAt(json, afterWhitespace)
            if c == "]" then Ok((JsonArray(newAcc), afterWhitespace + 1))
            else if c == "," then parseArrayItems(json, skipWhitespace(json, afterWhitespace + 1), newAcc)
            else Err("Expected ',' or ']' at position " + toString(afterWhitespace))
        }
    }
}

fn parseObject(json: String, pos: Int): Result<(JsonValue, Int), String> = {
    // Skip opening brace and whitespace
    let startPos = skipWhitespace(json, pos + 1)

    if String.charAt(json, startPos) == "}" then Ok((JsonObject([]), startPos + 1))
    else parseObjectFields(json, startPos, [])
}

fn parseObjectFields(json: String, pos: Int, acc: List<(String, JsonValue)>): Result<(JsonValue, Int), String> = {
    // Parse key
    match parseString(json, pos) {
        Err(e) => Err(e),
        Ok((keyValue, afterKey)) => {
            match keyValue {
                JsonString(key) => {
                    let colonPos = skipWhitespace(json, afterKey)
                    if String.charAt(json, colonPos) != ":" then
                        Err("Expected ':' at position " + toString(colonPos))
                    else {
                        let valuePos = skipWhitespace(json, colonPos + 1)
                        match parseValue(json, valuePos) {
                            Err(e) => Err(e),
                            Ok((value, afterValue)) => {
                                let newAcc = List.concat(acc, [(key, value)])
                                let afterWhitespace = skipWhitespace(json, afterValue)
                                let c = String.charAt(json, afterWhitespace)
                                if c == "}" then Ok((JsonObject(newAcc), afterWhitespace + 1))
                                else if c == "," then parseObjectFields(json, skipWhitespace(json, afterWhitespace + 1), newAcc)
                                else Err("Expected ',' or '}' at position " + toString(afterWhitespace))
                            }
                        }
                    }
                },
                _ => Err("Expected string key at position " + toString(pos))
            }
        }
    }
}

fn skipWhitespace(json: String, pos: Int): Int = {
    let c = String.charAt(json, pos)
    if c == " " || c == "\n" || c == "\r" || c == "\t" then
        skipWhitespace(json, pos + 1)
    else pos
}

fn isDigit(c: String): Bool =
    c == "0" || c == "1" || c == "2" || c == "3" || c == "4" ||
    c == "5" || c == "6" || c == "7" || c == "8" || c == "9"

// ============================================================
// Codec Type (for automatic serialization)
// ============================================================

// A codec can both encode and decode a type
type Codec<A> = {
    encode: fn(A): JsonValue,
    decode: fn(JsonValue): Result<A, String>
}

// Create a codec from encode/decode functions
pub fn codec<A>(
    enc: fn(A): JsonValue,
    dec: fn(JsonValue): Result<A, String>
): Codec<A> =
    { encode: enc, decode: dec }

// Built-in codecs
pub fn stringCodec(): Codec<String> =
    codec(
        encodeString,
        fn(json: JsonValue): Result<String, String> => match json {
            JsonString(s) => Ok(s),
            _ => Err("Expected string")
        }
    )

pub fn intCodec(): Codec<Int> =
    codec(
        encodeInt,
        fn(json: JsonValue): Result<Int, String> => match json {
            JsonInt(n) => Ok(n),
            _ => Err("Expected int")
        }
    )

pub fn boolCodec(): Codec<Bool> =
    codec(
        encodeBool,
        fn(json: JsonValue): Result<Bool, String> => match json {
            JsonBool(b) => Ok(b),
            _ => Err("Expected bool")
        }
    )

pub fn listCodec<A>(itemCodec: Codec<A>): Codec<List<A>> =
    codec(
        fn(items: List<A>): JsonValue => encodeList(items, itemCodec.encode),
        fn(json: JsonValue): Result<List<A>, String> => match json {
            JsonArray(items) => decodeAll(items, itemCodec.decode),
            _ => Err("Expected array")
        }
    )

fn decodeAll<A>(items: List<JsonValue>, decode: fn(JsonValue): Result<A, String>): Result<List<A>, String> = {
    match List.head(items) {
        None => Ok([]),
        Some(item) => match decode(item) {
            Err(e) => Err(e),
            Ok(decoded) => match decodeAll(Option.getOrElse(List.tail(items), []), decode) {
                Err(e) => Err(e),
                Ok(rest) => Ok(List.concat([decoded], rest))
            }
        }
    }
}

pub fn optionCodec<A>(itemCodec: Codec<A>): Codec<Option<A>> =
    codec(
        fn(opt: Option<A>): JsonValue => encodeOption(opt, itemCodec.encode),
        fn(json: JsonValue): Result<Option<A>, String> => match json {
            JsonNull => Ok(None),
            _ => itemCodec.decode(json).mapResult(fn(a: A): Option<A> => Some(a))
        }
    )

// ============================================================
// Helper for Result.mapResult
// ============================================================

fn mapResult<A, B>(result: Result<A, String>, f: fn(A): B): Result<B, String> =
    match result {
        Ok(a) => Ok(f(a)),
        Err(e) => Err(e)
    }