type LiteralValue =
    | { kind: 'int', value: number }
    | { kind: 'float', value: number }
    | { kind: 'string', value: string };

// Literals and Variables

export type Literal = {
    kind: 'literal'
    value: LiteralValue
    line?: number
    column?: number
    start?: number
}

export type Variable = {
    kind: 'variable'
    name: string
    line?: number
    column?: number
    start?: number
}

export type Constructor = {
    kind: 'constructor'
    name: string
    line?: number
    column?: number
    start?: number
}

// Functions

export type Lambda = {
    kind: 'lambda'
    params: string[]
    body: AST
    line?: number
    column?: number
    start?: number
}

export type Apply = {
    kind: 'apply'
    func: AST
    args: AST[]
    line?: number
    column?: number
    start?: number
}

// Bindings

export type Let = {
    kind: 'let'
    name: string
    value: AST
    body: AST
    line?: number
    column?: number
    start?: number
}

// Matching

export type Match = {
    kind: 'match'
    expr: AST
    cases: MatchCase[]
    line?: number
    column?: number
    start?: number
}

export type MatchCase = {
    pattern: Pattern
    result: AST
    line?: number
    column?: number
    start?: number
}

export type Pattern =
    | { kind: 'wildcard' }
    | { kind: 'var', name: string }
    | { kind: 'literal', value: number | string }
    | { kind: 'constructor', name: string, args: Pattern[] }
    | { kind: 'list', elements: Pattern[] }
    | { kind: 'list-spread', head: Pattern[], spread: string }
    | { kind: 'record', fields: { [key: string]: Pattern } }

// Data Structures

export type ListSpread = {
    kind: 'list-spread'
    spread: AST;
    line?: number;
    column?: number;
    start?: number;
}

export type List = {
    kind: 'list'
    elements: (AST | ListSpread)[]
    line?: number
    column?: number
    start?: number
}

export type Record = {
    kind: 'record'
    entries: Array<
    | { kind: 'field', key: string, value: AST }
    | { kind: 'spread', expr: AST }
    >
    // fields: { [key: string]: AST }
    line?: number
    column?: number
    start?: number
}

export type RecordAccess = {
    kind: 'record-access'
    record: AST
    field: string
    line?: number
    column?: number
    start?: number
}

export type RecordUpdate = {
    kind: 'record-update'
    record: AST
    updates: { [key: string]: AST }
    line?: number
    column?: number
    start?: number
}

// Top-level constructs

export type Definition = {
    kind: 'definition'
    name: string
    body: AST
    line?: number
    column?: number
    start?: number
}

export type Rebind = {
    kind: 'rebind'
    target: AST
    value: AST
    line?: number
    column?: number
    start?: number
}

export type AST =
    | Literal
    | Variable
    | Constructor
    | Lambda
    | Apply
    | Let
    | Match
    | Record
    | RecordAccess
    | RecordUpdate
    | List
    | ListSpread
    | Definition
    | Rebind

export function prettyPrint(ast: AST, indent = 0): string {
    const i = '  '.repeat(indent);

    switch (ast.kind) {
        case 'literal': {
            const val = ast.value;
            if (val.kind === 'string') {
                return `${i}"${val.value}"`;
            }
            return `${i}${val.value}`;
        }

        case 'variable':
        case 'constructor':
            return ast.name;

        case 'apply':
            const func = prettyPrint(ast.func, 0);
            const args = ast.args.map(a => {
                const printed = prettyPrint(a, 0);
                if (a.kind === 'lambda' || a.kind === 'match' || a.kind === 'let' || a.kind === 'rebind') {
                    return `(${printed})`;
                }
                return printed;
            }).join(' ');

            return `(${func} ${args})`

        case 'let':
            return `${ast.name} = ${prettyPrint(ast.value, indent + 1)};\n${i}${prettyPrint(ast.body, indent)}`

        case 'list':
            const elems = ast.elements.map(e => prettyPrint(e, 0)).join(', ');
            return `[${elems}]`;

        case 'record':
            const parts = ast.entries.map(entry =>
                entry.kind === 'spread'
                    ? `...${prettyPrint(entry.expr, )}`
                    : `${entry.key} = ${prettyPrint(entry.value, 0)}`
            );
            return `{ ${parts.join(', ')} }`;

        case 'lambda': {
            const params = ast.params.join(' ');
            const body = prettyPrint(ast.body, indent + 1);
            const isComplex = ast.body.kind === 'match' || ast.body.kind === 'let';
            if (isComplex) {
                return `${params} \\\n${body}`

            }
            return `${params} \\ ${body}`
        }

        case 'record-access':
            return `${prettyPrint(ast.record, 0)}.${ast.field}`;

        case 'record-update': {
            const updates = Object.entries(ast.updates)
                .map(([k, v]) => `${k} = ${prettyPrint(v, 0)}`)
                .join(', ');
            return `${prettyPrint(ast.record, 0)}.{ ${updates} }`
        }

        case 'match':
            const expr = prettyPrint(ast.expr, 0);
            const cases = ast.cases
                .map(c => `${i}| ${prettyPrintPattern(c.pattern)} \\ ${prettyPrint(c.result, indent + 1)}`)
                .join('\n');
            return `${expr}\n${cases}`;

        case 'rebind':
            return `${prettyPrint(ast.target, 0)} := ${prettyPrint(ast.value, 0)}`;

        case 'list-spread':
            return `...${prettyPrint(ast.spread, 0)}`;

        case 'definition':
            return `${ast.name} = ${prettyPrint(ast.body, indent)};`;

        default:
            return `Unknown AST kind: ${i}${(ast as any).kind}`
    }
}

function prettyPrintPattern(pattern: Pattern): string {
    switch (pattern.kind) {
        case 'wildcard':
            return '_';

        case 'var':
            return pattern.name;

        case 'literal':
            return JSON.stringify(pattern.value);

        case 'constructor':
            if (pattern.args.length === 0) {
                return pattern.name;
            }
            const args = pattern.args.map(prettyPrintPattern).join(' ');
            return `(${pattern.name} ${args})`;

        case 'list':
            const elems = pattern.elements.map(prettyPrintPattern).join(', ');
            return `[${elems}]`;

        case 'list-spread':
            const head = pattern.head.map(prettyPrintPattern).join(', ');
            return head.length > 0
                ? `[${head}, ...${pattern.spread}]`
                : `[...${pattern.spread}]`;

        case 'record':
            const fields = Object.entries(pattern.fields)
                .map(([k, p]) => `${k} = ${prettyPrintPattern(p)}`)
                .join(', ');

            return `{${fields}}`;
        default:
            return `Unknown AST kind: ${(pattern as any).kind}`
    }
}