using StellaOps.Policy;
using System.Collections.Immutable;

namespace StellaOps.PolicyDsl;

/// <summary>
/// Parses policy DSL source code into an AST.
/// </summary>
public sealed class PolicyParser
{
    private readonly ImmutableArray<DslToken> tokens;
    private readonly List<PolicyIssue> diagnostics = new();
    private int position;

    private PolicyParser(ImmutableArray<DslToken> tokens)
    {
        this.tokens = tokens;
    }

    public static PolicyParseResult Parse(string source)
    {
        if (source is null)
        {
            throw new ArgumentNullException(nameof(source));
        }

        var tokenization = DslTokenizer.Tokenize(source);
        var parser = new PolicyParser(tokenization.Tokens);
        var document = parser.ParseDocument();
        var allDiagnostics = tokenization.Diagnostics.AddRange(parser.diagnostics).ToImmutableArray();
        return new PolicyParseResult(document, allDiagnostics);
    }

    private PolicyDocumentNode? ParseDocument()
    {
        if (!Match(TokenKind.KeywordPolicy))
        {
            diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.MissingPolicyHeader, "Expected 'policy' declaration.", "policy"));
            return null;
        }

        var nameToken = Consume(TokenKind.StringLiteral, "Policy name must be a string literal.", "policy.name");
        var name = nameToken.Value as string ?? nameToken.Text;

        Consume(TokenKind.KeywordSyntax, "Expected 'syntax' declaration.", "policy.syntax");
        var syntaxToken = Consume(TokenKind.StringLiteral, "Policy syntax must be a string literal.", "policy.syntax.value");
        var syntax = syntaxToken.Value as string ?? syntaxToken.Text;

        Consume(TokenKind.LeftBrace, "Expected '{' to start policy body.", "policy.body");

        var metadataBuilder = ImmutableDictionary.CreateBuilder<string, PolicyLiteralValue>(StringComparer.Ordinal);
        var settingsBuilder = ImmutableDictionary.CreateBuilder<string, PolicyLiteralValue>(StringComparer.Ordinal);
        var profiles = ImmutableArray.CreateBuilder<PolicyProfileNode>();
        var rules = ImmutableArray.CreateBuilder<PolicyRuleNode>();

        while (!Check(TokenKind.RightBrace) && !IsAtEnd)
        {
            if (Match(TokenKind.KeywordMetadata))
            {
                foreach (var kvp in ParseKeyValueBlock("policy.metadata"))
                {
                    metadataBuilder[kvp.Key] = kvp.Value;
                }

                continue;
            }

            if (Match(TokenKind.KeywordSettings))
            {
                foreach (var kvp in ParseKeyValueBlock("policy.settings"))
                {
                    settingsBuilder[kvp.Key] = kvp.Value;
                }

                continue;
            }

            if (Match(TokenKind.KeywordProfile))
            {
                var profile = ParseProfile();
                if (profile is not null)
                {
                    profiles.Add(profile);
                }

                continue;
            }

            if (Match(TokenKind.KeywordRule))
            {
                var rule = ParseRule();
                if (rule is not null)
                {
                    rules.Add(rule);
                }

                continue;
            }

            diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.UnexpectedSection, $"Unexpected token '{Current.Text}' in policy body.", "policy.body"));
            Advance();
        }

        var close = Consume(TokenKind.RightBrace, "Expected '}' to close policy definition.", "policy");

        if (!string.Equals(syntax, "stella-dsl@1", StringComparison.Ordinal))
        {
            diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.UnsupportedSyntaxVersion, $"Unsupported syntax '{syntax}'.", "policy.syntax"));
        }

        var span = new SourceSpan(tokens[0].Span.Start, close.Span.End);
        return new PolicyDocumentNode(
            name,
            syntax,
            metadataBuilder.ToImmutable(),
            profiles.ToImmutable(),
            settingsBuilder.ToImmutable(),
            rules.ToImmutable(),
            span);
    }

    private PolicyProfileNode? ParseProfile()
    {
        var nameToken = Consume(TokenKind.Identifier, "Profile requires a name.", "policy.profile");
        var name = nameToken.Text;
        Consume(TokenKind.LeftBrace, "Expected '{' after profile declaration.", $"policy.profile.{name}");

        var start = nameToken.Span.Start;
        var depth = 1;
        while (depth > 0 && !IsAtEnd)
        {
            if (Match(TokenKind.LeftBrace))
            {
                depth++;
            }
            else if (Match(TokenKind.RightBrace))
            {
                depth--;
            }
            else
            {
                Advance();
            }
        }

        var close = Previous;
        return new PolicyProfileNode(
            name,
            ImmutableArray<PolicyProfileItemNode>.Empty,
            new SourceSpan(start, close.Span.End));
    }

    private PolicyRuleNode? ParseRule()
    {
        var nameToken = Consume(TokenKind.Identifier, "Rule requires a name.", "policy.rule");
        var name = nameToken.Text;

        var priority = 0;
        if (Match(TokenKind.KeywordPriority))
        {
            var priorityToken = Consume(TokenKind.NumberLiteral, "Priority must be numeric.", $"policy.rule.{name}");
            if (priorityToken.Value is decimal dec)
            {
                priority = (int)Math.Round(dec, MidpointRounding.AwayFromZero);
            }
        }

        Consume(TokenKind.LeftBrace, "Expected '{' to start rule.", $"policy.rule.{name}");
        Consume(TokenKind.KeywordWhen, "Rule requires a 'when' clause.", $"policy.rule.{name}");
        var when = ParseExpression();

        Consume(TokenKind.KeywordThen, "Rule requires a 'then' clause.", $"policy.rule.{name}");
        var thenActions = ParseActions(name, "then");

        var elseActions = ImmutableArray<PolicyActionNode>.Empty;
        if (Match(TokenKind.KeywordElse))
        {
            elseActions = ParseActions(name, "else");
        }

        string? because = null;
        if (Match(TokenKind.KeywordBecause))
        {
            var becauseToken = Consume(TokenKind.StringLiteral, "Because clause must be string.", $"policy.rule.{name}.because");
            because = becauseToken.Value as string ?? becauseToken.Text;
        }

        var close = Consume(TokenKind.RightBrace, "Expected '}' to close rule.", $"policy.rule.{name}");

        if (because is null)
        {
            diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.MissingBecauseClause, $"Rule '{name}' missing 'because' clause.", $"policy.rule.{name}"));
        }

        return new PolicyRuleNode(name, priority, when, thenActions, elseActions, because, new SourceSpan(nameToken.Span.Start, close.Span.End));
    }

    private ImmutableArray<PolicyActionNode> ParseActions(string ruleName, string clause)
    {
        var actions = ImmutableArray.CreateBuilder<PolicyActionNode>();
        while (!Check(TokenKind.RightBrace) && !Check(TokenKind.KeywordElse) && !Check(TokenKind.KeywordBecause) && !IsAtEnd)
        {
            if (Check(TokenKind.Identifier))
            {
                actions.Add(ParseAssignmentAction(ruleName, clause));
                continue;
            }

            if (Match(TokenKind.KeywordAnnotate))
            {
                actions.Add(ParseAnnotateAction(ruleName, clause));
                continue;
            }

            if (Match(TokenKind.KeywordWarn))
            {
                actions.Add(ParseWarnAction());
                continue;
            }

            if (Match(TokenKind.KeywordEscalate))
            {
                actions.Add(ParseEscalateAction());
                continue;
            }

            if (Match(TokenKind.KeywordRequireVex))
            {
                actions.Add(ParseRequireVexAction(ruleName, clause));
                continue;
            }

            if (Match(TokenKind.KeywordIgnore))
            {
                actions.Add(ParseIgnoreAction(ruleName, clause));
                continue;
            }

            if (Match(TokenKind.KeywordDefer))
            {
                actions.Add(ParseDeferAction(ruleName, clause));
                continue;
            }

            diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.InvalidAction, $"Unexpected token '{Current.Text}' in {clause} actions.", $"policy.rule.{ruleName}.{clause}"));
            Advance();
        }

        return actions.ToImmutable();
    }

    private PolicyActionNode ParseAssignmentAction(string ruleName, string clause)
    {
        var segments = ImmutableArray.CreateBuilder<string>();
        var first = Consume(TokenKind.Identifier, "Assignment target must start with identifier.", $"policy.rule.{ruleName}.{clause}");
        segments.Add(first.Text);
        while (Match(TokenKind.Dot))
        {
            segments.Add(Consume(TokenKind.Identifier, "Expected identifier after '.'.", $"policy.rule.{ruleName}.{clause}").Text);
        }

        Consume(TokenKind.Define, "Expected ':=' in action.", $"policy.rule.{ruleName}.{clause}");
        var value = ParseExpression();
        Match(TokenKind.Semicolon);
        return new PolicyAssignmentActionNode(new PolicyReference(segments.ToImmutable(), new SourceSpan(first.Span.Start, value.Span.End)), value, new SourceSpan(first.Span.Start, value.Span.End));
    }

    private PolicyActionNode ParseAnnotateAction(string ruleName, string clause)
    {
        var reference = ParseReference($"policy.rule.{ruleName}.{clause}.annotate");
        Consume(TokenKind.Define, "Expected ':=' in annotate action.", $"policy.rule.{ruleName}.{clause}.annotate");
        var value = ParseExpression();
        Match(TokenKind.Semicolon);
        return new PolicyAnnotateActionNode(reference, value, new SourceSpan(reference.Span.Start, value.Span.End));
    }

    private PolicyActionNode ParseWarnAction()
    {
        PolicyExpression? message = null;
        if (Match(TokenKind.KeywordMessage))
        {
            message = ParseExpression();
        }

        Match(TokenKind.Semicolon);
        var span = message?.Span ?? Previous.Span;
        return new PolicyWarnActionNode(message, span);
    }

    private PolicyActionNode ParseEscalateAction()
    {
        PolicyExpression? to = null;
        PolicyExpression? when = null;

        if (Match(TokenKind.KeywordTo))
        {
            to = ParseExpression();
        }

        if (Match(TokenKind.KeywordWhen))
        {
            when = ParseExpression();
        }

        Match(TokenKind.Semicolon);
        var end = when?.Span.End ?? to?.Span.End ?? Previous.Span.End;
        return new PolicyEscalateActionNode(to, when, new SourceSpan(Previous.Span.Start, end));
    }

    private PolicyActionNode ParseRequireVexAction(string ruleName, string clause)
    {
        Consume(TokenKind.LeftBrace, "Expected '{' after requireVex.", $"policy.rule.{ruleName}.{clause}.requireVex");
        var builder = ImmutableDictionary.CreateBuilder<string, PolicyExpression>(StringComparer.Ordinal);
        while (!Check(TokenKind.RightBrace) && !IsAtEnd)
        {
            var key = Consume(TokenKind.Identifier, "requireVex key must be identifier.", $"policy.rule.{ruleName}.{clause}.requireVex").Text;
            Consume(TokenKind.Assign, "Expected '=' in requireVex condition.", $"policy.rule.{ruleName}.{clause}.requireVex");
            builder[key] = ParseExpression();
            Match(TokenKind.Comma);
        }

        var close = Consume(TokenKind.RightBrace, "Expected '}' to close requireVex block.", $"policy.rule.{ruleName}.{clause}.requireVex");
        Match(TokenKind.Semicolon);
        return new PolicyRequireVexActionNode(builder.ToImmutable(), new SourceSpan(close.Span.Start, close.Span.End));
    }

    private PolicyActionNode ParseIgnoreAction(string ruleName, string clause)
    {
        PolicyExpression? until = null;
        string? because = null;
        if (Match(TokenKind.KeywordUntil))
        {
            until = ParseExpression();
        }

        if (Match(TokenKind.KeywordBecause))
        {
            var becauseToken = Consume(TokenKind.StringLiteral, "Ignore 'because' must be string.", $"policy.rule.{ruleName}.{clause}.ignore");
            because = becauseToken.Value as string ?? becauseToken.Text;
        }

        Match(TokenKind.Semicolon);
        return new PolicyIgnoreActionNode(until, because, new SourceSpan(Previous.Span.Start, (until?.Span.End ?? Previous.Span.End)));
    }

    private PolicyActionNode ParseDeferAction(string ruleName, string clause)
    {
        PolicyExpression? until = null;
        if (Match(TokenKind.KeywordUntil))
        {
            until = ParseExpression();
        }

        Match(TokenKind.Semicolon);
        return new PolicyDeferActionNode(until, new SourceSpan(Previous.Span.Start, (until?.Span.End ?? Previous.Span.End)));
    }

    private PolicyReference ParseReference(string path)
    {
        var segments = ImmutableArray.CreateBuilder<string>();
        var first = Consume(TokenKind.Identifier, "Expected identifier.", path);
        segments.Add(first.Text);
        while (Match(TokenKind.Dot))
        {
            segments.Add(Consume(TokenKind.Identifier, "Expected identifier after '.'.", path).Text);
        }

        return new PolicyReference(segments.ToImmutable(), first.Span);
    }

    private Dictionary<string, PolicyLiteralValue> ParseKeyValueBlock(string path)
    {
        Consume(TokenKind.LeftBrace, "Expected '{'.", path);
        var entries = new Dictionary<string, PolicyLiteralValue>(StringComparer.Ordinal);
        while (!Check(TokenKind.RightBrace) && !IsAtEnd)
        {
            var key = Consume(TokenKind.Identifier, "Expected identifier.", path).Text;
            Consume(TokenKind.Assign, "Expected '='.", path);
            entries[key] = ParseLiteralValue(path);
            Match(TokenKind.Semicolon);
        }

        Consume(TokenKind.RightBrace, "Expected '}'.", path);
        return entries;
    }

    private PolicyLiteralValue ParseLiteralValue(string path)
    {
        if (Match(TokenKind.StringLiteral))
        {
            return new PolicyStringLiteral(Previous.Value as string ?? Previous.Text, Previous.Span);
        }

        if (Match(TokenKind.NumberLiteral))
        {
            return new PolicyNumberLiteral(Previous.Value is decimal dec ? dec : 0m, Previous.Span);
        }

        if (Match(TokenKind.BooleanLiteral))
        {
            return new PolicyBooleanLiteral(Previous.Value is bool b && b, Previous.Span);
        }

        if (Match(TokenKind.LeftBracket))
        {
            var start = Previous.Span.Start;
            var items = ImmutableArray.CreateBuilder<PolicyLiteralValue>();
            while (!Check(TokenKind.RightBracket) && !IsAtEnd)
            {
                items.Add(ParseLiteralValue(path));
                Match(TokenKind.Comma);
            }

            var close = Consume(TokenKind.RightBracket, "Expected ']' in list literal.", path);
            return new PolicyListLiteral(items.ToImmutable(), new SourceSpan(start, close.Span.End));
        }

        diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.InvalidLiteral, "Invalid literal.", path));
        return new PolicyStringLiteral(string.Empty, Current.Span);
    }

    private PolicyExpression ParseExpression() => ParseOr();

    private PolicyExpression ParseOr()
    {
        var expr = ParseAnd();
        while (Match(TokenKind.KeywordOr))
        {
            var right = ParseAnd();
            expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.Or, right, new SourceSpan(expr.Span.Start, right.Span.End));
        }

        return expr;
    }

    private PolicyExpression ParseAnd()
    {
        var expr = ParseEquality();
        while (Match(TokenKind.KeywordAnd))
        {
            var right = ParseEquality();
            expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.And, right, new SourceSpan(expr.Span.Start, right.Span.End));
        }

        return expr;
    }

    private PolicyExpression ParseEquality()
    {
        var expr = ParseUnary();
        while (true)
        {
            if (Match(TokenKind.EqualEqual))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.Equal, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else if (Match(TokenKind.NotEqual))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.NotEqual, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else if (Match(TokenKind.KeywordIn))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.In, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else if (Match(TokenKind.KeywordNot))
            {
                if (Match(TokenKind.KeywordIn))
                {
                    var right = ParseUnary();
                    expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.NotIn, right, new SourceSpan(expr.Span.Start, right.Span.End));
                }
                else
                {
                    diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.UnexpectedToken, "Expected 'in' after 'not'.", "expression.not"));
                }
            }
            else if (Match(TokenKind.LessThan))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.LessThan, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else if (Match(TokenKind.LessThanOrEqual))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.LessThanOrEqual, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else if (Match(TokenKind.GreaterThan))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.GreaterThan, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else if (Match(TokenKind.GreaterThanOrEqual))
            {
                var right = ParseUnary();
                expr = new PolicyBinaryExpression(expr, PolicyBinaryOperator.GreaterThanOrEqual, right, new SourceSpan(expr.Span.Start, right.Span.End));
            }
            else
            {
                break;
            }
        }

        return expr;
    }

    private PolicyExpression ParseUnary()
    {
        if (Match(TokenKind.KeywordNot))
        {
            var operand = ParseUnary();
            return new PolicyUnaryExpression(PolicyUnaryOperator.Not, operand, new SourceSpan(Previous.Span.Start, operand.Span.End));
        }

        return ParsePrimary();
    }

    private PolicyExpression ParsePrimary()
    {
        if (Match(TokenKind.StringLiteral))
        {
            return new PolicyLiteralExpression(Previous.Value as string ?? Previous.Text, Previous.Span);
        }

        if (Match(TokenKind.NumberLiteral))
        {
            return new PolicyLiteralExpression(Previous.Value ?? 0m, Previous.Span);
        }

        if (Match(TokenKind.BooleanLiteral))
        {
            return new PolicyLiteralExpression(Previous.Value ?? false, Previous.Span);
        }

        if (Match(TokenKind.LeftBracket))
        {
            var start = Previous.Span.Start;
            var items = ImmutableArray.CreateBuilder<PolicyExpression>();
            while (!Check(TokenKind.RightBracket) && !IsAtEnd)
            {
                items.Add(ParseExpression());
                Match(TokenKind.Comma);
            }

            var close = Consume(TokenKind.RightBracket, "Expected ']' to close list expression.", "expression.list");
            return new PolicyListExpression(items.ToImmutable(), new SourceSpan(start, close.Span.End));
        }

        if (Match(TokenKind.LeftParen))
        {
            var expr = ParseExpression();
            Consume(TokenKind.RightParen, "Expected ')' to close grouped expression.", "expression.group");
            return expr;
        }

        if (Match(TokenKind.Identifier))
        {
            return ParseIdentifierExpression(Previous);
        }

        if (Match(TokenKind.KeywordEnv))
        {
            return ParseIdentifierExpression(Previous);
        }

        diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.UnexpectedToken, $"Unexpected token '{Current.Text}' in expression.", "expression"));
        var bad = Advance();
        return new PolicyLiteralExpression(null, bad.Span);
    }

    private PolicyExpression ParseIdentifierExpression(DslToken identifier)
    {
        PolicyExpression expr = new PolicyIdentifierExpression(identifier.Text, identifier.Span);
        while (true)
        {
            if (Match(TokenKind.Dot))
            {
                var member = ConsumeIdentifier("Expected identifier after '.'.", "expression.member");
                expr = new PolicyMemberAccessExpression(expr, member.Text, new SourceSpan(expr.Span.Start, member.Span.End));
                continue;
            }

            if (Match(TokenKind.LeftParen))
            {
                var args = ImmutableArray.CreateBuilder<PolicyExpression>();
                if (!Check(TokenKind.RightParen))
                {
                    do
                    {
                        args.Add(ParseExpression());
                    }
                    while (Match(TokenKind.Comma));
                }

                var close = Consume(TokenKind.RightParen, "Expected ')' to close invocation.", "expression.call");
                expr = new PolicyInvocationExpression(expr, args.ToImmutable(), new SourceSpan(expr.Span.Start, close.Span.End));
                continue;
            }

            if (Match(TokenKind.LeftBracket))
            {
                var indexExpr = ParseExpression();
                var close = Consume(TokenKind.RightBracket, "Expected ']' to close indexer.", "expression.indexer");
                expr = new PolicyIndexerExpression(expr, indexExpr, new SourceSpan(expr.Span.Start, close.Span.End));
                continue;
            }

            break;
        }

        return expr;
    }

    private DslToken ConsumeIdentifier(string message, string path)
    {
        if (Check(TokenKind.Identifier) || IsKeywordIdentifier(Current.Kind))
        {
            return Advance();
        }

        diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.UnexpectedToken, message, path));
        return Advance();
    }

    private static bool IsKeywordIdentifier(TokenKind kind) =>
        kind == TokenKind.KeywordSource;

    private bool Match(TokenKind kind)
    {
        if (Check(kind))
        {
            Advance();
            return true;
        }

        return false;
    }

    private bool Check(TokenKind kind) => !IsAtEnd && Current.Kind == kind;

    private DslToken Consume(TokenKind kind, string message, string path)
    {
        if (Check(kind))
        {
            return Advance();
        }

        diagnostics.Add(PolicyIssue.Error(DiagnosticCodes.UnexpectedToken, message, path));
        return Advance();
    }

    private DslToken Advance()
    {
        if (!IsAtEnd)
        {
            position++;
        }

        return tokens[position - 1];
    }

    private bool IsAtEnd => Current.Kind == TokenKind.EndOfFile;

    private DslToken Current => tokens[position];

    private DslToken Previous => tokens[position - 1];
}

/// <summary>
/// Result of parsing a policy DSL source.
/// </summary>
public readonly record struct PolicyParseResult(
    PolicyDocumentNode? Document,
    ImmutableArray<PolicyIssue> Diagnostics);