using Microsoft.SqlServer.TransactSql.ScriptDom;

namespace SQLLinter.Infrastructure.Diagram;

public enum BpmnNodeType
{
    Start,
    End,
    Task,
    Gateway,
    Subprocess,
}

public class BpmnNode
{
    public string Id { get; set; } = string.Empty;
    public string Label { get; set; } = string.Empty;
    public BpmnNodeType Type { get; set; }
}

public class BpmnEdge
{
    public string From { get; set; } = string.Empty;
    public string To { get; set; } = string.Empty;
    public string Label { get; set; } = string.Empty;
    public bool Dashed { get; set; }
}

public class BpmnProcess
{
    public string Id { get; set; } = string.Empty;
    public string Name { get; set; } = string.Empty;
    public List<BpmnNode> Nodes { get; set; } = new();
    public List<BpmnEdge> Edges { get; set; } = new();
}

public class BpmnDiagram
{
    public BpmnProcess Main { get; set; } = new BpmnProcess { Id = "Main", Name = "Main" };
    public List<BpmnProcess> Subprocesses { get; set; } = new();
}

public static class BpmnBuilder
{
    public static BpmnDiagram Build(TSqlFragment fragment)
    {
        var visitor = new BpmnVisitor();
        fragment.Accept(visitor);
        return visitor.Diagram;
    }

    private class BpmnVisitor : TSqlFragmentVisitor
    {
        public BpmnDiagram Diagram { get; } = new BpmnDiagram();

        private Dictionary<string, string> _lastNodeByProcess = new();
        private BpmnProcess _currentProcess;
        private int _nodeCounter = 0;

        public BpmnVisitor()
        {
            _currentProcess = Diagram.Main;
            // create a Start node for main
            var start = new BpmnNode { Id = NewId(), Label = "Start", Type = BpmnNodeType.Start };
            Diagram.Main.Nodes.Add(start);
            _lastNodeByProcess[Diagram.Main.Id] = start.Id;
        }

        private string NewId() => "n" + (_nodeCounter++);

        private void AddNode(BpmnProcess proc, BpmnNode node)
        {
            if (string.IsNullOrEmpty(node.Id)) node.Id = NewId();
            proc.Nodes.Add(node);

            if (_lastNodeByProcess.TryGetValue(proc.Id, out var last) && !string.IsNullOrEmpty(last))
            {
                proc.Edges.Add(new BpmnEdge { From = last, To = node.Id });
            }
            _lastNodeByProcess[proc.Id] = node.Id;
        }

        private void AddNodeToCurrent(string label, BpmnNodeType type)
        {
            var node = new BpmnNode { Id = NewId(), Label = label, Type = type };
            AddNode(_currentProcess, node);
        }

        public override void Visit(TSqlScript node)
        {
            _currentProcess = Diagram.Main;
            if (!_lastNodeByProcess.ContainsKey(_currentProcess.Id)) _lastNodeByProcess[_currentProcess.Id] = null;
            base.Visit(node);

            // add End node for main process
            var end = new BpmnNode { Id = NewId(), Label = "End", Type = BpmnNodeType.End };
            AddNode(Diagram.Main, end);
        }

        public override void Visit(CreateProcedureStatement node)
        {
            // AST-based extraction of procedure name
            var name = FindSchemaObjectName(node) ?? ("proc" + Diagram.Subprocesses.Count);
            var pid = SanitizeId(name);
            var proc = new BpmnProcess { Id = pid, Name = name };

            // synthetic pool/label node
            var poolNode = new BpmnNode { Id = pid + "_pool", Label = name, Type = BpmnNodeType.Subprocess };
            proc.Nodes.Add(poolNode);
            // create start node for subprocess
            var start = new BpmnNode { Id = pid + "_start", Label = "Start", Type = BpmnNodeType.Start };
            proc.Nodes.Add(start);
            // set last node pointer to start
            _lastNodeByProcess[proc.Id] = start.Id;

            Diagram.Subprocesses.Add(proc);

            var old = _currentProcess;
            _currentProcess = proc;
            if (!_lastNodeByProcess.ContainsKey(proc.Id)) _lastNodeByProcess[proc.Id] = start.Id;

            base.Visit(node);

            // after visiting procedure body add End node
            var end = new BpmnNode { Id = pid + "_end", Label = "End", Type = BpmnNodeType.End };
            AddNode(proc, end);

            _currentProcess = old;
        }

        public override void Visit(ExecuteStatement node)
        {
            // AST-based extraction for called proc
            var called = FindProcedureNameFromExecute(node);
            var label = called != null ? $"EXEC {called}" : node.ToString();

            // task in current process
            var task = new BpmnNode { Id = NewId(), Label = label, Type = BpmnNodeType.Task };
            AddNode(_currentProcess, task);

            // if calling known subprocess, add dashed edge from this node to subprocess start
            if (called != null)
            {
                var proc = Diagram.Subprocesses.FirstOrDefault(p => string.Equals(p.Name, called, StringComparison.OrdinalIgnoreCase) || string.Equals(p.Id, SanitizeId(called), StringComparison.OrdinalIgnoreCase));
                if (proc != null)
                {
                    // add dashed edge in main process (from this task to subprocess pool node)
                    Diagram.Main.Edges.Add(new BpmnEdge { From = task.Id, To = proc.Nodes.First().Id, Dashed = true });
                }
            }

            base.Visit(node);
        }

        public override void Visit(IfStatement node)
        {
            var gid = new BpmnNode { Id = NewId(), Label = "IF", Type = BpmnNodeType.Gateway };
            AddNode(_currentProcess, gid);

            // then / else placeholders
            if (node.ThenStatement != null)
            {
                var thenNode = new BpmnNode { Id = NewId(), Label = "THEN", Type = BpmnNodeType.Task };
                AddNode(_currentProcess, thenNode);
                _currentProcess.Edges.Add(new BpmnEdge { From = gid.Id, To = thenNode.Id });
            }

            if (node.ElseStatement != null)
            {
                var elseNode = new BpmnNode { Id = NewId(), Label = "ELSE", Type = BpmnNodeType.Task };
                AddNode(_currentProcess, elseNode);
                _currentProcess.Edges.Add(new BpmnEdge { From = gid.Id, To = elseNode.Id });
            }

            base.Visit(node);
        }

        public override void Visit(SelectStatement node)
        {
            var label = "SELECT";
            if (node.QueryExpression is QuerySpecification qs && qs.SelectElements != null)
            {
                label = $"SELECT ({qs.SelectElements.Count})";
            }
            AddNodeToCurrent(label, BpmnNodeType.Task);
            base.Visit(node);
        }

        public override void Visit(InsertStatement node)
        {
            var target = FindTargetNameForDml(node);
            var label = "INSERT" + (string.IsNullOrEmpty(target) ? string.Empty : " -> " + target);
            AddNodeToCurrent(label, BpmnNodeType.Task);
            base.Visit(node);
        }

        public override void Visit(UpdateStatement node)
        {
            var target = FindTargetNameForDml(node);
            var label = "UPDATE" + (string.IsNullOrEmpty(target) ? string.Empty : " -> " + target);
            AddNodeToCurrent(label, BpmnNodeType.Task);
            base.Visit(node);
        }

        public override void Visit(DeleteStatement node)
        {
            var target = FindTargetNameForDml(node);
            var label = "DELETE" + (string.IsNullOrEmpty(target) ? string.Empty : " -> " + target);
            AddNodeToCurrent(label, BpmnNodeType.Task);
            base.Visit(node);
        }

        // AST helpers
        private static string? FindSchemaObjectName(TSqlFragment fragment)
        {
            var finder = new SchemaNameFinder();
            fragment.Accept(finder);
            return finder.FoundName;
        }

        private static string? FindProcedureNameFromExecute(ExecuteStatement node)
        {
            // Prefer ProcedureReference inside the statement
            var finder = new ProcedureReferenceFinder();
            node.Accept(finder);
            return finder.FoundName;
        }

        private static string? FindTargetNameForDml(TSqlFragment fragment)
        {
            var finder = new TargetTableFinder();
            fragment.Accept(finder);
            return finder.FoundName;
        }

        private class SchemaNameFinder : TSqlFragmentVisitor
        {
            public string? FoundName { get; private set; }
            public override void Visit(SchemaObjectName node)
            {
                if (FoundName == null)
                {
                    try
                    {
                        if (node.Identifiers != null && node.Identifiers.Count > 0)
                        {
                            FoundName = string.Join('.', node.Identifiers.Select(id => id.Value));
                        }
                        else
                        {
                            FoundName = node.ToString();
                        }
                    }
                    catch
                    {
                        FoundName = node.ToString();
                    }
                }
                base.Visit(node);
            }

            public override void Visit(ProcedureReferenceName node)
            {
                if (FoundName == null)
                {
                    try
                    {
                        if (node.ProcedureReference != null && node.ProcedureReference.Name != null)
                        {
                            // procedure reference name may contain identifiers
                            FoundName = node.ProcedureReference.Name.ToString();
                        }
                        else
                        {
                            FoundName = node.ToString();
                        }
                    }
                    catch
                    {
                        FoundName = node.ToString();
                    }
                }
                base.Visit(node);
            }
        }

        private class ProcedureReferenceFinder : TSqlFragmentVisitor
        {
            public string? FoundName { get; private set; }
            public override void Visit(ProcedureReference node)
            {
                if (FoundName == null)
                {
                    try
                    {
                        if (node.Name != null)
                        {
                            if (node.Name.Identifiers != null && node.Name.Identifiers.Count > 0)
                                FoundName = string.Join('.', node.Name.Identifiers.Select(id => id.Value));
                            else
                                FoundName = node.Name.ToString();
                        }
                    }
                    catch
                    {
                        FoundName = node.ToString();
                    }
                }
                base.Visit(node);
            }

            public override void Visit(ExecuteSpecification node)
            {
                base.Visit(node);
            }
        }

        private class TargetTableFinder : TSqlFragmentVisitor
        {
            public string? FoundName { get; private set; }
            private int _bestDepth = int.MaxValue;
            private int _depth = 0;

            public override void Visit(TSqlFragment node)
            {
                _depth++;
                base.Visit(node);
                _depth--;
            }

            public override void Visit(NamedTableReference node)
            {
                if (node == null) return;
                int depth = _depth;
                if (depth < _bestDepth)
                {
                    _bestDepth = depth;
                    try
                    {
                        if (node.SchemaObject != null && node.SchemaObject.Identifiers != null && node.SchemaObject.Identifiers.Count > 0)
                            FoundName = string.Join('.', node.SchemaObject.Identifiers.Select(id => id.Value));
                        else
                            FoundName = node.SchemaObject?.ToString() ?? node.ToString();

                        if (node.Alias != null && !string.IsNullOrWhiteSpace(node.Alias.Value))
                            FoundName += " AS " + node.Alias.Value;
                    }
                    catch
                    {
                        FoundName = node.SchemaObject?.ToString() ?? node.ToString();
                    }
                }

                base.Visit(node);
            }

            public override void Visit(VariableTableReference node)
            {
                if (node == null) return;
                int depth = _depth;
                if (depth < _bestDepth)
                {
                    _bestDepth = depth;
                    FoundName = node.Variable?.ToString();
                }
                base.Visit(node);
            }

            public override void Visit(QueryDerivedTable node)
            {
                if (node == null) return;
                int depth = _depth;
                if (depth < _bestDepth && node.Alias != null && !string.IsNullOrWhiteSpace(node.Alias.Value))
                {
                    _bestDepth = depth;
                    FoundName = node.Alias.Value + " (derived)";
                }

                // still traverse inside
                base.Visit(node);
            }

            public override void Visit(CommonTableExpression node)
            {
                // CTE definitions: prefer not to pick CTE as target unless no other table found
                base.Visit(node);
            }
        }

        private static string SanitizeId(string s)
        {
            if (string.IsNullOrEmpty(s)) return "id";
            var sb = new System.Text.StringBuilder();
            foreach (var ch in s)
            {
                if (char.IsLetterOrDigit(ch)) sb.Append(ch);
                else sb.Append('_');
            }
            return sb.ToString();
        }
    }
}