using System.Reflection.Metadata;
using System.Text;
using System.Text.RegularExpressions;
using Antlr4.Runtime;
using meowlang.parser.antlr;

namespace meowlang.parser;

public class InvalidEscapeSequenceException : Exception
{
    public InvalidEscapeSequenceException(string escapeSequence) : base($"invalid escape sequence `{escapeSequence}`")
    {
    }
}

public class MeowVisitorException : Exception
{
    public IToken Token { get; }

    public MeowVisitorException(IToken token, string message) : base(message)
    {
        Token = token;
    }
}

public class InvalidImportPathException : MeowVisitorException
{
    public InvalidImportPathException(IToken context, string message) : base(context, message)
    {
    }
}

public class Visitor : MeowBaseVisitor<Model>
{
    public override Model VisitFile(MeowParser.FileContext context)
    {
        var imports = context.importStatement().Select(x => new ImportVisitor().Visit(x)).ToList();
        var declarations = context.declaration().Select(x => new DeclarationVisitor().Visit(x)).ToList();

        return new Model(context.GetSpan(), imports, declarations);
    }
}

public class ImportVisitor : MeowBaseVisitor<ImportModel>
{
    private static readonly string PathSegmentPattern = "[a-z][a-z0-9_]*";

    private static readonly Regex ProjectRegex = new Regex("^[a-z][a-z0-9_]*$");

    private static readonly Regex PathRegexNoProject =
        new Regex($"^({PathSegmentPattern}|\\.|\\.\\.)(/{PathSegmentPattern})*$");

    private static readonly Regex PathRegex = new Regex($"^{PathSegmentPattern}(/{PathSegmentPattern})*$");

    public override ImportModel VisitImportStatement(MeowParser.ImportStatementContext context)
    {
        var path = context.importpath.Text[1..^1].Unescape();

        string? project = null;

        if (path.Contains(':'))
        {
            var parts = path.Split(':');
            project = parts[0];
            if (!ProjectRegex.IsMatch(project))
            {
                throw new InvalidImportPathException(context.importpath, "malformed project name");
            }

            path = parts[1];
            if (!PathRegex.IsMatch(path))
            {
                throw new InvalidImportPathException(context.importpath, "malformed path");
            }
        }
        else
        {
            if (!PathRegexNoProject.IsMatch(path))
            {
                throw new InvalidImportPathException(context.importpath, "malformed path");
            }
        }

        var alias = context.importname?.Text;

        return new ImportModel(context.GetSpan(), project, path, alias);
    }
}

public class DeclarationVisitor : MeowBaseVisitor<DeclarationModel>
{
    public override DeclarationModel VisitFunction(MeowParser.FunctionContext context)
    {
        var attributes = context.attribute().Select(x => new AttributeVisitor().Visit(x)).ToList();
        var pub = context.pub != null;
        var name = context.name.Text;

        var typeParameters =
            context.genericParameters()?._name.Select(x => new TypeParameterModel(x.GetSpan(), x.Text)).ToList() ??
            new List<TypeParameterModel>();
        var functionParameters = context.functionParameters().functionParameter()
            .Select(x => new FunctionParameterVisitor().Visit(x)).ToList();
        var functionReturns = context.functionReturn()?.typeReference().Select(x => new TypeReferenceVisitor().Visit(x))
            .ToList() ?? new List<TypeReferenceModel>();
        var constraints = context.constraint()?.Select(x => new FunctionConstraintVisitor().Visit(x)).ToList() ??
                          new List<ConstraintModel>();

        return new FunctionModel(context.GetSpan(), attributes, pub, name, typeParameters, functionParameters,
            functionReturns, constraints,
            new FunctionBodyModel(context.functionBody().GetSpan())); // TODO: visit function body
    }

    public override ConstraintDeclarationModel VisitConstraintDeclaration(
        MeowParser.ConstraintDeclarationContext context)
    {
        var name = context.name.Text;
        var typeNames = context._types.Select(x => x.Text).ToList();
        var rules = context.constraintRule().Select(x => new ConstraintRuleVisitor().Visit(x)).ToList();
        return new ConstraintDeclarationModel(context.GetSpan(), name, typeNames, rules);
    }

    public override DeclarationModel VisitStructDeclaration(MeowParser.StructDeclarationContext context)
    {
        var name = context.name.Text;
        var typeParameters = context.genericParameters()?._name.Select(x => x.Text).ToList() ?? new List<string>();
        var constraints = context.constraint()?.Select(x => new FunctionConstraintVisitor().Visit(x)).ToList() ??
                          new List<ConstraintModel>();
        var structModel = new TypeReferenceVisitor().Visit(context.structType()) as StructTypeModel;
        return new StructDeclarationModel(context.GetSpan(), name, typeParameters, constraints, structModel!);
    }

    public override DeclarationModel VisitEnumDeclaration(MeowParser.EnumDeclarationContext context)
    {
        var name = context.name.Text;
        var typeParameters = context.genericParameters()?._name.Select(x => x.Text).ToList() ?? new List<string>();
        var constraints = context.constraint()?.Select(x => new FunctionConstraintVisitor().Visit(x)).ToList() ??
                          new List<ConstraintModel>();

        var enumModel = new TypeReferenceVisitor().Visit(context.enumType()) as EnumTypeModel;
        return new EnumDeclarationModel(context.GetSpan(), name, typeParameters, constraints, enumModel!);
    }

    public override DeclarationModel VisitTupleDeclaration(MeowParser.TupleDeclarationContext context)
    {
        var name = context.name.Text;
        var typeParameters = context.genericParameters()?._name.Select(x => x.Text).ToList() ?? new List<string>();
        var constraints = context.constraint()?.Select(x => new FunctionConstraintVisitor().Visit(x)).ToList() ??
                          new List<ConstraintModel>();
        var tupleModel = new TypeReferenceVisitor().Visit(context.tupleType()) as TupleTypeModel;
        return new TupleDeclarationModel(context.GetSpan(), name, typeParameters, constraints, tupleModel!);
    }

    public override DeclarationModel VisitTypeAlias(MeowParser.TypeAliasContext context)
    {
        var name = context.name.Text;
        var typeParameters = context.genericParameters()?._name.Select(x => x.Text).ToList() ?? new List<string>();
        var type = new TypeReferenceVisitor().Visit(context.typeReference());
        return new TypeAliasModel(context.GetSpan(), name, typeParameters, type);
    }
}

public class AttributeVisitor : MeowBaseVisitor<AttributeModel>
{
    public override AttributeModel VisitAttribute(MeowParser.AttributeContext context)
    {
        var name = context.name.Text;
        var values = context.attributeValue().Select(x => new AttributeValueVisitor().Visit(x)).ToList();
        return new AttributeModel(context.GetSpan(), name, values);
    }
}

public class AttributeValueVisitor : MeowBaseVisitor<AttributeValueModel>
{
    public override AttributeValueModel VisitAttributeValue(MeowParser.AttributeValueContext context)
    {
        if (context.id != null)
        {
            return new IdentifierAttributeValueModel(context.GetSpan(), context.id.Text);
        }

        var literal = new LiteralVisitor().Visit(context.literal());
        return new LiteralAttributeValueModel(context.GetSpan(), literal);
    }
}

public class FunctionParameterVisitor : MeowBaseVisitor<FunctionParameterModel>
{
    public override FunctionParameterModel VisitFunctionParameter(MeowParser.FunctionParameterContext context)
    {
        var type = new TypeReferenceVisitor().Visit(context.typeReference());
        var name = context.name.Text;
        return new FunctionParameterModel(context.GetSpan(), type, name);
    }
}

public class TypeReferenceVisitor : MeowBaseVisitor<TypeReferenceModel>
{
    public override TypeReferenceModel VisitTypeReference(MeowParser.TypeReferenceContext context)
    {
        var type = Visit(context.nonPointerTypeReference());
        if (context.pointer != null)
        {
            return new PointerTypeReferenceModel(context.GetSpan(), type);
        }

        return type;
    }

    public override TypeReferenceModel VisitTypeName(MeowParser.TypeNameContext context)
    {
        var import = context.importName?.Text;
        var nestedName = context._name.Select(x => x.Text).ToList();

        var typeParameters = context._genericType.Select(Visit).ToList();

        return new TypeNameModel(context.GetSpan(), import, nestedName[0], typeParameters, nestedName.Skip(1).ToList());
    }

    public override TypeReferenceModel VisitArrayType(MeowParser.ArrayTypeContext context)
    {
        var type = Visit(context.typeReference());
        if (context.length == null)
        {
            return new SliceTypeModel(context.GetSpan(), type);
        }

        return new ArrayTypeModel(context.GetSpan(), type, context.length.Text);
    }

    public override TypeReferenceModel VisitTupleType(MeowParser.TupleTypeContext context)
    {
        var types = context.typeReference().Select(Visit).ToList();
        return new TupleTypeModel(context.GetSpan(), types);
    }

    public override TypeReferenceModel VisitStructType(MeowParser.StructTypeContext context)
    {
        var members = context.structMember().Select(x => new StructMemberVisitor().Visit(x)).ToList();
        return new StructTypeModel(context.GetSpan(), members);
    }

    public override TypeReferenceModel VisitEnumType(MeowParser.EnumTypeContext context)
    {
        var members = context.enumMember().Select(x => new EnumMemberVisitor().Visit(x)).ToList();
        return new EnumTypeModel(context.GetSpan(), members);
    }
}

public class StructMemberVisitor : MeowBaseVisitor<StructMemberModel>
{
    public override StructMemberModel VisitStructMember(MeowParser.StructMemberContext context)
    {
        var type = new TypeReferenceVisitor().Visit(context.typeReference());
        var name = context.name.Text;
        return new StructMemberModel(context.GetSpan(), type, name);
    }
}

public class EnumMemberVisitor : MeowBaseVisitor<EnumMemberModel>
{
    public override EnumMemberModel VisitEnumMember(MeowParser.EnumMemberContext context)
    {
        var name = context.name.Text;

        TypeReferenceModel? type = null;
        var typeRef = context.typeReference();
        if (typeRef != null)
        {
            type = new TypeReferenceVisitor().Visit(typeRef);
        }

        return new EnumMemberModel(context.GetSpan(), name, type);
    }
}

public class FunctionConstraintVisitor : MeowBaseVisitor<ConstraintModel>
{
    public override ConstraintModel VisitConstraint(MeowParser.ConstraintContext context)
    {
        var name = context.name.Text;
        var typeNames = context._typenames.Select(x => x.Text).ToList();
        return new ConstraintModel(context.GetSpan(), name, typeNames);
    }
}

public class LiteralVisitor : MeowBaseVisitor<LiteralModel>
{
    public override LiteralModel VisitStringLiteral(MeowParser.StringLiteralContext context)
    {
        var value = context.val.Text[1..^1].Unescape();
        return new StringLiteralModel(context.GetSpan(), value);
    }

    public override LiteralModel VisitBoolLiteral(MeowParser.BoolLiteralContext context)
    {
        var text = context.val.Text;
        var value = text == "true";
        return new BoolLiteralModel(context.GetSpan(), value);
    }

    public override LiteralModel VisitNumericLiteral(MeowParser.NumericLiteralContext context)
    {
        return new NumericLiteralModel(context.GetSpan(), context.val.Text);
    }

    public override LiteralModel VisitCharacterLiteral(MeowParser.CharacterLiteralContext context)
    {
        var value = context.val.Text[1..^1].Unescape();
        return new CharacterLiteralModel(context.GetSpan(), value);
    }
}

public class ConstraintRuleVisitor : MeowBaseVisitor<ConstraintRuleModel>
{
    public override ConstraintRuleModel VisitEmbedConstraintRule(MeowParser.EmbedConstraintRuleContext context)
    {
        var name = context.name.Text;
        var types = context._types.Select(x => new TypeReferenceVisitor().Visit(x)).ToList();
        return new EmbedConstraintRuleModel(context.GetSpan(), name, types);
    }

    public override ConstraintRuleModel VisitFunctionConstraintRule(MeowParser.FunctionConstraintRuleContext context)
    {
        var name = context.name.Text;
        var parameters = context.constraintFunctionParameters().typeReference()
            .Select(x => new TypeReferenceVisitor().Visit(x)).ToList();
        var returns =
            context.functionReturn()?.typeReference().Select(x => new TypeReferenceVisitor().Visit(x)).ToList() ??
            new List<TypeReferenceModel>();
        return new FunctionConstraintRuleModel(context.GetSpan(), name, parameters, returns);
    }
}

public record Model(Span span, List<ImportModel> Imports, List<DeclarationModel> Declarations) : ModelBase(span);

public record ImportModel(Span span, string? Project, string Path, string? Alias) : ModelBase(span);

public abstract record DeclarationModel(Span span) : ModelBase(span);

public record FunctionModel(Span span, List<AttributeModel> Attributes, bool Pub, string Name,
    List<TypeParameterModel> TypeParameters, List<FunctionParameterModel> Parameters, List<TypeReferenceModel> Returns,
    List<ConstraintModel> Constraints, FunctionBodyModel Body) : DeclarationModel(span);

public record AttributeModel(Span span, string Name, List<AttributeValueModel> Values) : ModelBase(span);

public abstract record AttributeValueModel(Span span) : ModelBase(span);

public record IdentifierAttributeValueModel(Span span, string Name) : AttributeValueModel(span);

public record LiteralAttributeValueModel(Span span, LiteralModel Literal) : AttributeValueModel(span);

public record TypeParameterModel(Span span, string Name) : ModelBase(span);

public record FunctionParameterModel(Span span, TypeReferenceModel Type, string Name) : ModelBase(span);

public abstract record TypeReferenceModel(Span span) : ModelBase(span);

public record PointerTypeReferenceModel(Span span, TypeReferenceModel Type) : TypeReferenceModel(span);

public record TypeNameModel(Span span, string? ImportName, string Name, List<TypeReferenceModel> TypeParameters,
    List<string> Nested) : TypeReferenceModel(span);

public record ArrayTypeModel(Span span, TypeReferenceModel Type, string Length) : TypeReferenceModel(span);

public record SliceTypeModel(Span span, TypeReferenceModel Type) : TypeReferenceModel(span);

public record TupleTypeModel(Span span, List<TypeReferenceModel> Types) : TypeReferenceModel(span);

public record StructTypeModel(Span span, List<StructMemberModel> Members) : TypeReferenceModel(span);

public record StructMemberModel(Span span, TypeReferenceModel Type, string Name) : ModelBase(span);

public record EnumTypeModel(Span span, List<EnumMemberModel> Members) : TypeReferenceModel(span);

public record EnumMemberModel(Span span, string Name, TypeReferenceModel? Type) : ModelBase(span);

public record ConstraintModel(Span span, string Name, List<string> TypeNames) : ModelBase(span);

public record FunctionBodyModel(Span span) : ModelBase(span);

public abstract record LiteralModel(Span span) : ModelBase(span);

public record StringLiteralModel(Span span, string Value) : LiteralModel(span);

public record CharacterLiteralModel(Span span, string Value) : LiteralModel(span);

public record NumericLiteralModel(Span span, string StringValue) : LiteralModel(span);

public record BoolLiteralModel(Span span, bool Value) : LiteralModel(span);

public record ConstraintDeclarationModel
    (Span span, string Name, List<string> TypeNames, List<ConstraintRuleModel> Rules) : DeclarationModel(span);

public abstract record ConstraintRuleModel(Span span) : ModelBase(span);

public record EmbedConstraintRuleModel
    (Span span, string Name, List<TypeReferenceModel> Types) : ConstraintRuleModel(span);

public record FunctionConstraintRuleModel(Span span, string Name, List<TypeReferenceModel> Parameters,
    List<TypeReferenceModel> Returns) : ConstraintRuleModel(span);

public record StructDeclarationModel(Span span, string Name, List<string> TypeParameters,
    List<ConstraintModel> Constraints, StructTypeModel Struct) : DeclarationModel(span);

public record EnumDeclarationModel(Span span, string Name, List<string> TypeParameters,
    List<ConstraintModel> Constraints, EnumTypeModel Enum) : DeclarationModel(span);

public record TupleDeclarationModel(Span span, string Name, List<string> TypeParameters,
    List<ConstraintModel> Constraints, TupleTypeModel Tuple) : DeclarationModel(span);

public record TypeAliasModel
    (Span span, string Name, List<string> TypeParameters, TypeReferenceModel Type) : DeclarationModel(span);