import type { AST, Pattern } from './ast';
import type { Env } from './env';
import type { Value } from './types';
import { RuntimeError } from './error';
import { recordDependency } from './store';

export function evaluate(ast: AST, env: Env, source: string): Value {
    switch (ast.kind) {
        case 'literal':
            return ast.value;

        case 'variable': {
            const val = env.get(ast.name);

            if (val === undefined)
                throw RuntimeError(`Unknown variable: ${ast.name}`, ast.line, ast.column, source);

            recordDependency(ast.name);

            return val;
        }

        case 'list': {
            const elements: Value[] = [];

            for (const item of ast.elements) {
                // Spread
                if ('spread' in item) {
                    const spreadValue = evaluate(item.spread, env, source);

                    if (spreadValue.kind !== 'list')
                        throw RuntimeError('can only spread lists', ast.line, ast.column, source);

                    elements.push(...spreadValue.elements);
                } else {
                    elements.push(evaluate(item, env, source));
                }
            }

            return { kind: 'list', elements };
        }

        case 'record': {
            const fields: { [key: string]: Value } = {};
            const fieldMeta: { [key: string]: { body: AST, dependencies: Set<string> } } = {};
            const recordEnv = new Map(env);
            const fieldNames = Object.keys(ast.fields);

            for (const [k, fieldAst] of Object.entries(ast.fields)) {
                // Track which siblings are accessed
                const deps = new Set<string>();

                const trackingEnv = new Map(recordEnv);
                const originalGet = trackingEnv.get.bind(trackingEnv);
                trackingEnv.get = (name: string) => {
                    if (fieldNames.includes(name) && name !== k) {
                        deps.add(name);
                    }
                    return originalGet(name);
                }

                const value = evaluate(fieldAst, trackingEnv, source);

                fields[k] = value;
                fieldMeta[k] = { body: fieldAst, dependencies: deps };
                recordEnv.set(k, value);
            }

            return { kind: 'record', fields, fieldMeta };
        }

        case 'record-access': {
            const record = evaluate(ast.record, env, source);

            if (record.kind !== 'record')
                throw RuntimeError('Not a record', ast.line, ast.column, source);

            const value = record.fields[ast.field];

            if (value === undefined) {
                throw RuntimeError(`Field ${ast.field} not found`, ast.line, ast.column, source);
            }

            return value;
        }

        case 'record-update': {
            const record = evaluate(ast.record, env, source);

            if (record.kind !== 'record')
                throw RuntimeError('Not a record', ast.line, ast.column, source);

            const newFields: { [key: string]: Value } = { ...record.fields };

            for (const [field, expr] of Object.entries(ast.updates)) {
                newFields[field] = evaluate(expr, env, source);
            }

            return { kind: 'record', fields: newFields };
        }

        case 'constructor':
            return {
                kind: 'constructor',
                name: ast.name,
                args: []
            };

        case 'let': {
            const newEnv = new Map(env);

            const val = evaluate(ast.value, newEnv, source);

            // Don't bind _
            if (ast.name !== '_') {
                newEnv.set(ast.name, val);
            }

            newEnv.set(ast.name, val);

            return evaluate(ast.body, newEnv, source);
        }

        case 'lambda':
            return {
                kind: 'closure',
                params: ast.params,
                body: ast.body,
                env
            }

        case 'apply': {
            const func = evaluate(ast.func, env, source);
            const argValues = ast.args.map(arg => evaluate(arg, env, source));

            // Native functions
            if (func.kind === 'native') {
                // Exact args
                if (argValues.length === func.arity) {
                    return func.fn(...argValues);
                }

                // Partial application
                if (argValues.length < func.arity) {
                    const capturedArgs = argValues;

                    return {
                        kind: 'native',
                        name: func.name,
                        arity: func.arity - argValues.length,
                        fn: (...restArgs: Value[]) => {
                            return func.fn(...capturedArgs, ...restArgs);

                        }
                    };
                }

                throw RuntimeError(`Function expects ${func.arity} args, but got ${argValues.length}`, ast.line, ast.column, source);
            }

            // Constructor application
            if (func.kind === 'constructor') {
                const argValues = ast.args.map(arg => evaluate(arg, env, source));
                return {
                    kind: 'constructor',
                    name: func.name,
                    args: [...func.args, ...argValues]
                };
            }

            if (func.kind !== 'closure')
                throw RuntimeError('Not a function', ast.line, ast.column, source);

            // Too few args (Currying)
            if (argValues.length < func.params.length) {
                // Bind the params we have
                const newEnv = new Map(func.env);

                for (let i = 0; i < argValues.length; i++) {
                    newEnv.set(func.params[i], argValues[i]);
                }

                return {
                    kind: 'closure',
                    params: func.params.slice(argValues.length),
                    body: func.body,
                    env: newEnv
                };
            }

            // Too many args
            if (argValues.length > func.params.length)
                throw RuntimeError('Too many arguments', ast.line, ast.column, source);

            // Exact number of args
            const callEnv = new Map(func.env);
            for (let i = 0; i < argValues.length; i++) {
                callEnv.set(func.params[i], argValues[i]);
            }

            return evaluate(func.body, callEnv, source);
        }

        case 'match': {
            const value = evaluate(ast.expr, env, source);

            for (const matchCase of ast.cases) {
                const bindings = matchPattern(value, matchCase.pattern);

                if (bindings !== null) {
                    const newEnv = new Map(env);
                    for (const [name, val] of Object.entries(bindings)) {
                        newEnv.set(name, val);
                    }
                    return evaluate(matchCase.result, newEnv, source);
                }
            }

            throw RuntimeError('Non-exhaustive pattern match', ast.line, ast.column, source);
        }

        case 'rebind': {
            const value = evaluate(ast.value, env, source);

            if (ast.target.kind === 'variable') {
                const name = ast.target.name;
                return {
                    kind: 'constructor',
                    name: 'Rebind',
                    args: [{ kind: 'string', value: name }, value]
                };
            }

            if (ast.target.kind === 'record-access') {
                let current: AST = ast.target;
                const path: string[] = [];

                while (current.kind === 'record-access') {
                    path.unshift(current.field);
                    current = current.record;
                }

                if (current.kind !== 'variable')
                    throw RuntimeError('Rebind target must be a variable or field access', ast.line, ast.column, source);

                const rootName = current.name;
                const rootValue = env.get(rootName);

                if (!rootValue)
                    throw RuntimeError(`Unknown variable: ${rootName}`, ast.line, ast.column, source);

                return {
                    kind: 'constructor',
                    name: 'Rebind',
                    args: [
                        { kind: 'string', value: rootName },
                        { kind: 'list', elements: path.map(p => ({ kind: 'string', value: p })) },
                        value
                    ]
                };
            }

            throw RuntimeError('Rebind target must be a variable or field access', ast.line, ast.column, source);
        }

        default:
            throw RuntimeError('Syntax Error', ast.line, ast.column, source);
    }
}

type Bindings = { [key: string]: Value };

function matchPattern(value: Value, pattern: Pattern): Bindings | null {
    switch (pattern.kind) {
        case 'wildcard':
            return {};

        case 'var':
            return { [pattern.name]: value };

        case 'literal':
            if (value.kind === 'int' || value.kind === 'float' || value.kind === 'string') {
                if (value.value === pattern.value) {
                    return {};
                }
            }
            return null;

        case 'constructor': {
            if (value.kind !== 'constructor') return null;
            if (value.name !== pattern.name) return null;
            if (value.args.length !== pattern.args.length) return null;

            const bindings: Bindings = {};
            for (let i = 0; i < pattern.args.length; i++) {
                const argBindings = matchPattern(value.args[i], pattern.args[i]);
                if (argBindings === null) return null;
                Object.assign(bindings, argBindings);
            }
            return bindings;
        }

        case 'list': {
            if (value.kind !== 'list') return null;
            if (value.elements.length !== pattern.elements.length) return null;

            const bindings: Bindings = {};
            for (let i = 0; i < pattern.elements.length; i++) {
                const elemBindings = matchPattern(value.elements[i], pattern.elements[i]);
                if (elemBindings === null) return null;
                Object.assign(bindings, elemBindings);
            }
            return bindings;
        }

        case 'list-spread': {
            if (value.kind !== 'list') return null;
            if (value.elements.length < pattern.head.length) return null;

            const bindings: Bindings = {};
            for (let i = 0; i < pattern.head.length; i++) {
                const elemBindings = matchPattern(value.elements[i], pattern.head[i]);
                if (elemBindings === null) return null;
                Object.assign(bindings, elemBindings);
            }

            const rest = value.elements.slice(pattern.head.length);
            bindings[pattern.spread] = { kind: 'list', elements: rest };

            return bindings;
        }

        case 'record': {
            if (value.kind !== 'record') return null;

            const bindings: Bindings = {};
            for (const [fieldName, fieldPattern] of Object.entries(pattern.fields)) {
                const fieldValue = value.fields[fieldName];
                if (fieldValue === undefined) return null;

                const fieldBindings = matchPattern(fieldValue, fieldPattern);
                if (fieldBindings === null) return null;
                Object.assign(bindings, fieldBindings);
            }
            return bindings;
        }
    }
}