using System;
using System.Collections.Immutable;
using System.Linq;
using Finn;
using Finn.AST;

using ScopeStack = System.Collections.Immutable.ImmutableStack<System.Collections.Immutable.ImmutableHashSet<string>>;
using Resolutions = System.Collections.Immutable.ImmutableDictionary<Finn.AST.Expr, int>;
using System.Collections.Generic;

class Resolver : IExprVisitor<ScopeStack, Resolutions>
{
	private class NameCollection
	{
		public ImmutableHashSet<string> Names { get; private set; } = ImmutableHashSet<string>.Empty;

		public void Add(Token name)
		{
			var value = (string)name.Literal!;
			if (Names.Contains(value))
			{
				Program.error(name.Start, $"Duplicate name declared in scope: {name.Lexeme}");
			}
			else
			{
				Names = Names.Add(value);
			}
		}
	}

	private class PatternNameCollector : IPatternVisitor<NameCollection, ValueTuple>
	{
		public readonly static PatternNameCollector Instance = new();
		private static readonly ValueTuple Void = ValueTuple.Create();

		private PatternNameCollector() { }

		public ValueTuple visitFieldPatternPattern(NameCollection names, FieldPattern field)
		{
			if (field.Pattern is null)
			{
				names.Add(field.Name);
			}
			else
			{
				field.Pattern.accept(names, this);
			}
			return Void;
		}

		public ValueTuple visitRecordPatternPattern(NameCollection names, RecordPattern record)
		{
			record.Rest?.accept(names, this);
			foreach (var field in record.Fields)
			{
				field.accept(names, this);
			}
			return Void;
		}

		public ValueTuple visitSimplePatternPattern(NameCollection names, SimplePattern simple)
		{
			if (simple.Identifier is not null)
			{
				names.Add(simple.Identifier);
			}
			return Void;
		}

		public ValueTuple visitVariantPatternPattern(NameCollection names, VariantPattern variant)
		{
			variant.Argument?.accept(names, this);
			return Void;
		}
	}

	private static readonly Resolutions EmptyResolutions = Resolutions.Empty.WithComparers(ReferenceEqualityComparer.Instance);
	public static readonly Resolver Instance = new Resolver();

	private Resolver() { }

	private static Resolutions AddRange(Resolutions r1, Resolutions r2) => r1.AddRange(r2);

	private Resolutions Resolve(ScopeStack scopes, Expr expr)
	{
		return expr.accept(scopes, this);
	}

	private Resolutions Resolve(ScopeStack scopes, ImmutableHashSet<string> newScope, Expr value)
	{
		return Resolve(scopes.Push(newScope), value);
	}

	public Resolutions visitBinaryExpr(ScopeStack scopes, Binary expr)
	{
		var left = Resolve(scopes, expr.Left);
		var right = Resolve(scopes, expr.Right);
		return left.AddRange(right);
	}

	public Resolutions visitCallExpr(ScopeStack scopes, Call expr)
	{
		return (from arg in expr.Arguments
				where arg is not null
				select Resolve(scopes, arg))
		.Aggregate(Resolve(scopes, expr.Left), AddRange);
	}

	public Resolutions visitGroupingExpr(ScopeStack scopes, Grouping expr)
	{
		return Resolve(scopes, expr.Expression);
	}

	public Resolutions visitIfExpr(ScopeStack scopes, If expr)
	{
		var condition = Resolve(scopes, expr.Condition);
		var _then = Resolve(scopes, expr.Then);
		var _else = Resolve(scopes, expr.Else);
		return condition.AddRange(_then).AddRange(_else);
	}

	public Resolutions visitIndexerExpr(ScopeStack scopes, Indexer expr)
	{
		var left = Resolve(scopes, expr.Left);
		var index = Resolve(scopes, expr.Index);
		return left.AddRange(index);
	}

	public Resolutions visitLetExpr(ScopeStack scopes, Let expr)
	{
		var names = new NameCollection();
		foreach (var binding in expr.Bindings)
		{
			switch (binding)
			{
				case FuncBinding fb:
					names.Add(fb.Name);
					break;
				case VarBinding vb:
					vb.Pattern.accept(names, PatternNameCollector.Instance);
					break;
				default:
					throw new Exception("wtf there are no other binding types");
			}
		}

		var newScopes = scopes.Push(names.Names);

		// TODO Put the static constructivity check here.
		//     For now, just check that a definition only uses names that come before it.
		//     Note that this will preclude using shadowed variables in definitions.

		return (from binding in expr.Bindings
				select (binding switch
				{
					FuncBinding fb => ResolveFunction(newScopes, fb),
					VarBinding vb => Resolve(newScopes, vb.Value),
					_ => throw new Exception("wtf there are no other binding types"),
				}))
				.Aggregate(Resolve(newScopes, expr.Body), AddRange);

	}

	private Resolutions ResolveFunction(ScopeStack scopes, FuncBinding fb)
	{
		NameCollection names = new();
		foreach (var param in fb.Params)
		{
			param.accept(names, PatternNameCollector.Instance);
		}
		return Resolve(scopes, names.Names, fb.Value);
	}

	public Resolutions visitListExpr(ScopeStack scopes, List expr)
	{
		return (from element in expr.Elements
				select Resolve(scopes, element))
				.Aggregate(EmptyResolutions, AddRange);
	}

	public Resolutions visitLiteralExpr(ScopeStack scopes, Literal expr)
	{
		return EmptyResolutions;
	}

	public Resolutions visitRecordExpr(ScopeStack scopes, Record expr)
	{
		var resolutions = (from extension in expr.Extensions
						   select Resolve(scopes, extension.Value))
						   .Aggregate(EmptyResolutions, AddRange);
		if (expr.Base is not null)
		{
			resolutions = resolutions.AddRange(Resolve(scopes, expr.Base.Value));
			resolutions = (from update in expr.Base.Updates
						   select Resolve(scopes, update.Value))
						  .Aggregate(resolutions, AddRange);
		}

		return resolutions;
	}

	public Resolutions visitSelectorExpr(ScopeStack scopes, Selector expr)
	{
		return Resolve(scopes, expr.Left);
	}

	public Resolutions visitSequenceExpr(ScopeStack scopes, Sequence expr)
	{
		return AddRange(Resolve(scopes, expr.Left), Resolve(scopes, expr.Right));
	}

	public Resolutions visitUnaryExpr(ScopeStack scopes, Unary expr)
	{
		return Resolve(scopes, expr.Right);
	}

	public Resolutions visitVariableExpr(ScopeStack scopes, Variable expr)
	{
		int depth = 0;
		while (!scopes.IsEmpty)
		{
			scopes = scopes.Pop(out var scope);
			if (scope.Contains((string)expr.Value.Literal!))
			{
				return EmptyResolutions.Add(expr, depth);
			}
			depth++;
		}
		// At this point, it must be a global. I don't like that we leave that resolution for runtime.
		// TODO Don't leave that resolution for runtime.
		// But I think that can wait until I make the resolver nameless.
		return EmptyResolutions.Add(expr, depth);
	}

	public Resolutions visitVariantExpr(ScopeStack scopes, Variant expr)
	{
		if (expr.Argument is not null)
		{
			return Resolve(scopes, expr.Argument);
		}
		return EmptyResolutions;
	}

	public Resolutions visitWhenExpr(ScopeStack scopes, When expr)
	{
		return (from _case in expr.Cases
				select ResolveCase(_case))
		.Aggregate(Resolve(scopes, expr.Head), AddRange);

		Resolutions ResolveCase(VarBinding _case)
		{
			var names = new NameCollection();
			_case.Pattern.accept(names, PatternNameCollector.Instance);
			return Resolve(scopes, names.Names, _case.Value);
		}
	}
}