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



class Resolver : IExprVisitor<ImmutableStack<ImmutableHashSet<string>>, ValueTuple>
{
	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 ValueTuple Void = ValueTuple.Create();
	private readonly Interpreter Interpreter;

	Resolver(Interpreter interpreter)
	{
		Interpreter = interpreter;
	}

	private void Resolve(ImmutableStack<ImmutableHashSet<string>> scopes, Expr expr)
	{
		expr.accept(scopes, this);
	}

	private void Resolve(ImmutableStack<ImmutableHashSet<string>> scopes, ImmutableHashSet<string> newScope, Expr value)
	{
		Resolve(scopes.Push(newScope), value);
	}

	public ValueTuple visitBinaryExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Binary expr)
	{
		Resolve(scopes, expr.Left);
		Resolve(scopes, expr.Right);
		return Void;
	}

	public ValueTuple visitCallExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Call expr)
	{
		Resolve(scopes, expr.Left);
		foreach (var arg in expr.Arguments)
		{
			if (arg is not null)
			{
				Resolve(scopes, arg);
			}
		}
		return Void;
	}

	public ValueTuple visitGroupingExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Grouping expr)
	{
		Resolve(scopes, expr.Expression);
		return Void;
	}

	public ValueTuple visitIfExpr(ImmutableStack<ImmutableHashSet<string>> scopes, If expr)
	{
		Resolve(scopes, expr.Condition);
		Resolve(scopes, expr.Then);
		Resolve(scopes, expr.Else);
		return Void;
	}

	public ValueTuple visitIndexerExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Indexer expr)
	{
		Resolve(scopes, expr.Left);
		Resolve(scopes, expr.Index);
		return Void;
	}

	public ValueTuple visitLetExpr(ImmutableStack<ImmutableHashSet<string>> 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.

		foreach (var binding in expr.Bindings)
		{
			switch (binding)
			{
				case FuncBinding fb:
					ResolveFunction(newScopes, fb);
					break;
				case VarBinding vb:
					Resolve(newScopes, vb.Value);
					break;
				default:
					throw new Exception("wtf there are no other binding types");
			}
		}
		Resolve(newScopes, expr.Body);

		return Void;
	}

	private void ResolveFunction(ImmutableStack<ImmutableHashSet<string>> scopes, FuncBinding fb)
	{
		NameCollection names = new();
		foreach (var param in fb.Params)
		{
			param.accept(names, PatternNameCollector.Instance);
		}
		Resolve(scopes, names.Names, fb.Value);
	}

	public ValueTuple visitListExpr(ImmutableStack<ImmutableHashSet<string>> scopes, List expr)
	{
		foreach (var element in expr.Elements)
		{
			Resolve(scopes, element);
		}
		return Void;
	}

	public ValueTuple visitLiteralExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Literal expr)
	{
		return Void;
	}

	public ValueTuple visitRecordExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Record expr)
	{
		if (expr.Base is not null)
		{
			Resolve(scopes, expr.Base.Value);
			foreach (var update in expr.Base.Updates)
			{
				if (update.Value is not null)
				{
					Resolve(scopes, update.Value);
				}
				else
				{
					// TODO Resolve update.Name like a var expression.
				}
			}
		}
		foreach (var extension in expr.Extensions)
		{
			if (extension.Value is not null)
			{
				Resolve(scopes, extension.Value);
			}
			else
			{
				// TODO Resolve extension.Name like a var expression.
			}
		}
		return Void;
	}

	public ValueTuple visitSelectorExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Selector expr)
	{
		Resolve(scopes, expr.Left);
		return Void;
	}

	public ValueTuple visitSequenceExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Sequence expr)
	{
		Resolve(scopes, expr.Left);
		Resolve(scopes, expr.Right);
		return Void;
	}

	public ValueTuple visitUnaryExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Unary expr)
	{
		Resolve(scopes, expr.Right);
		return Void;
	}

	public ValueTuple visitVariableExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Variable expr)
	{
		ResolveLocal(scopes, expr, expr.Value);
		return Void;
	}

	private void ResolveLocal(ImmutableStack<ImmutableHashSet<string>> scopes, Expr expr, Token name)
	{
		int depth = 0;
		while (!scopes.IsEmpty)
		{
			scopes = scopes.Pop(out var scope);
			if (scope.Contains((string)name.Literal!))
			{
				Interpreter.Resolve(expr, depth);
				return;
			}
			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.
	}

	public ValueTuple visitVariantExpr(ImmutableStack<ImmutableHashSet<string>> scopes, Variant expr)
	{
		if (expr.Argument is not null)
		{
			Resolve(scopes, expr.Argument);
		}
		return Void;
	}

	public ValueTuple visitWhenExpr(ImmutableStack<ImmutableHashSet<string>> scopes, When expr)
	{
		Resolve(scopes, expr.Head);
		foreach (var _case in expr.Cases)
		{
			var names = new NameCollection();
			_case.Pattern.accept(names, PatternNameCollector.Instance);
			Resolve(scopes, names.Names, _case.Value);
		}
		return Void;
	}
}