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
    annotation?: Annotation
}

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

export type TypeAST =
    | { kind: 'type-name', name: string } // Int, String, Bool
    | { kind: 'type-var', name: string } // a, b
    | { kind: 'type-function', param: TypeAST, result: TypeAST } // a \ b
    | { kind: 'type-apply', constructor: TypeAST, args: TypeAST[] } // List a, Maybe Int
    | { kind: 'type-record', fields: { name: string, type: TypeAST }[] } // { x: Int }

export type TypeConstructor = { name: string, args: TypeAST[] }

export type TypeDefinition = {
    kind: 'type-definition'
    name: string
    params: string[]
    constructors: TypeConstructor[]
    line?: number
    column?: number
    start?: number
}

export type Annotation = {
    constraints: Constraint[]
    type: TypeAST
}

export type Constraint = {
    className: string
    typeVar: string
}

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

const infixOps: { [key: string]: string } = {
    cat: '&', add: '+', sub: '-', mul: '*', div: '/', mod: '%',
    pow: '^', eq: '==', neq: '!=', gt: '>', lt: '<', gte: '>=', lte: '<='
};

const isInfix = (a: AST) => a.kind === 'apply' && a.func.kind === 'variable' && a.args.length === 2 && infixOps[a.func.name];

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') {
                const escaped = val.value
                    .replace(/\\/g, '\\\\')
                    .replace(/"/g, '\\"')
                    .replace(/\n/g, '\\n')
                    .replace(/\t/g, '\\t');
                return `"${escaped}"`;
            }
            return `${val.value}`;
        }

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

        case 'apply':
            // infix ops
            if (isInfix(ast)) {
                const wrapIfNeeded = (a: AST) => {
                    const printed = prettyPrint(a, indent);
                    if (a.kind === 'apply' || a.kind === 'lambda' || a.kind === 'match' || a.kind === 'let') {
                        return `(${printed})`;
                    }
                    return `${printed}`;

                }
                const left = wrapIfNeeded(ast.args[0]);
                const right = wrapIfNeeded(ast.args[1]);
                return `${left} ${infixOps[(ast.func as Variable).name]} ${right}`;
            }

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

            if (ast.func.kind === 'lambda' || ast.func.kind === 'match' || ast.func.kind === 'let') {
                return `(${func} ${args})`
            }
            return `${func} ${args}`

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

        case 'list': {
            const elems = ast.elements.map(e => prettyPrint(e, indent + 1))
            const oneLine = `[${elems.join(', ')}]`;
            if (oneLine.length <= 60 || elems.length <= 1) return oneLine;
            const inner = elems.map(e => `${'  '.repeat(indent + 1)}${e}`).join(',\n');
            return `[\n${inner}\n${i}]`;
        }

        case 'record':
            const parts = ast.entries.map(entry =>
                entry.kind === 'spread'
                    ? `...${prettyPrint(entry.expr, indent + 1)}`
                    : `${needsQuotes(entry.key) ? JSON.stringify(entry.key) : entry.key} = ${prettyPrint(entry.value, indent + 1)}`
            );
            const oneLine = `{ ${parts.join(', ')} }`;
            if (oneLine.length <= 60 || parts.length <= 1) return oneLine;
            // if (parts.length <= 1) return `{ ${parts.join(', ')} }`;
            const inner = parts.map(p => `${'  '.repeat(indent + 1)}${p}`).join(',\n');
            return `{\n${inner}\n${i}}`;

        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) {
                const ii = '  '.repeat(indent + 1);
                return `${params} \\\n${ii}${body}`
            }
            return `${params} \\ ${body}`
        }

        case 'record-access':
            const field = needsQuotes(ast.field) ? `${JSON.stringify(ast.field)}` : ast.field;
            return `${prettyPrint(ast.record, indent)}.${field}`;

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

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

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

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

        case 'definition':
            const ann = ast.annotation
                ? `${ast.name} : ${prettyPrintType(ast.annotation.type)};\n`
                : '';
            return `${ann}${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]) => `${needsQuotes(k) ? `${JSON.stringify(k)}` : k} = ${prettyPrintPattern(p)}`)
                .join(', ');

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

function prettyPrintType(type: TypeAST): string {
    switch (type.kind) {
        case 'type-name':
        case 'type-var':
            return type.name;
        case 'type-function':
            const param = type.param.kind === 'type-function'
                ? `(${prettyPrintType(type.param)})`
                : prettyPrintType(type.param);
            return `${param} \\ ${prettyPrintType(type.result)}`;
        case 'type-apply':
            const args = type.args.map(a =>
                a.kind === 'type-function' || a.kind === 'type-apply'
                    ? `(${prettyPrintType(a)})`
                    : prettyPrintType(a)
            ).join(' ');
            return `${prettyPrintType(type.constructor)} ${args}`;
        case 'type-record':
            const fields = type.fields
                .map(f => `${f.name} : ${prettyPrintType(f.type)}`)
                .join(', ');
            return `{ ${fields} }`;
    }
}

function prettyPrintTypeDefinition(td: TypeDefinition): string {
    const params = td.params.length > 0 ? ' ' + td.params.join(' ') : '';
    const ctors = td.constructors.map(c => {
        const args = c.args.map(a =>
            a.kind === 'type-function' || a.kind === 'type-apply'
                ? `(${prettyPrintType(a)})`
                : prettyPrintType(a)
        ).join(' ');
        return args ? `${c.name} ${args}` : c.name;
    }).join(' | ');
    return `${td.name}${params} = ${ctors};`;
}

function needsQuotes(key: string): boolean {
    return key === '_' || !/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(key);
}