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;

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

	private Resolutions Resolutions = EmptyResolutions;
	public bool TryResolve(Expr expr, out int distance)
	{
		return Resolutions.TryGetValue(expr, out distance);
	}

	public void Add(Expr expr)
	{
		var newResolutions = ResolverBuilder.Instance.Resolve(ScopeStack.Empty, expr);
		Resolutions = Resolutions.AddRange(newResolutions);
	}

	private class ResolverBuilder : 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;
			}
		}

		public static readonly ResolverBuilder Instance = new ResolverBuilder();

		private ResolverBuilder() { }

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

		public 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);
			}
		}
	}
}