package internal

import (
	"fmt"
	"go/ast"

	"golang.org/x/tools/go/ast/astutil"
)

// copyAST performs a deep copy of an AST. *ast.Ident identity will be
// preserved.
//
// This allows using astutil.Apply to rewrite an AST without modifying
// the original AST.
func copyAST(original ast.Node) ast.Node {
	// This function is necessarily long. No utility function exists to do this
	// clone, as most any attempt would need to have customization options, which
	// would need to be as expressive as Apply. A possibility to shorten the code
	// here would be to use reflection, but that trades clarity for shorter code.
	nodes := make(map[ast.Node]ast.Node)

	astutil.Apply(original, nil, func(cur *astutil.Cursor) bool {
		switch node := cur.Node().(type) {
		case nil:
			// No-op.
		case *ast.ArrayType:
			nodes[node] = &ast.ArrayType{
				Lbrack: node.Lbrack,
				Len:    exprFromMap(nodes, node.Len),
				Elt:    exprFromMap(nodes, node.Elt),
			}
		case *ast.AssignStmt:
			nodes[node] = &ast.AssignStmt{
				Lhs:    copyExprList(nodes, node.Lhs),
				TokPos: node.TokPos,
				Tok:    node.Tok,
				Rhs:    copyExprList(nodes, node.Rhs),
			}
		case *ast.BadDecl:
			nodes[node] = &ast.BadDecl{
				From: node.From,
				To:   node.To,
			}
		case *ast.BadExpr:
			nodes[node] = &ast.BadExpr{
				From: node.From,
				To:   node.To,
			}
		case *ast.BadStmt:
			nodes[node] = &ast.BadStmt{
				From: node.From,
				To:   node.To,
			}
		case *ast.BasicLit:
			nodes[node] = &ast.BasicLit{
				ValuePos: node.ValuePos,
				Kind:     node.Kind,
				Value:    node.Value,
			}
		case *ast.BinaryExpr:
			nodes[node] = &ast.BinaryExpr{
				X:     exprFromMap(nodes, node.X),
				OpPos: node.OpPos,
				Op:    node.Op,
				Y:     exprFromMap(nodes, node.Y),
			}
		case *ast.BlockStmt:
			nodes[node] = &ast.BlockStmt{
				Lbrace: node.Lbrace,
				List:   copyStmtList(nodes, node.List),
				Rbrace: node.Rbrace,
			}
		case *ast.BranchStmt:
			nodes[node] = &ast.BranchStmt{
				TokPos: node.TokPos,
				Tok:    node.Tok,
				Label:  identFromMap(nodes, node.Label),
			}
		case *ast.CallExpr:
			nodes[node] = &ast.CallExpr{
				Fun:      exprFromMap(nodes, node.Fun),
				Lparen:   node.Lparen,
				Args:     copyExprList(nodes, node.Args),
				Ellipsis: node.Ellipsis,
				Rparen:   node.Rparen,
			}
		case *ast.CaseClause:
			nodes[node] = &ast.CaseClause{
				Case:  node.Case,
				List:  copyExprList(nodes, node.List),
				Colon: node.Colon,
				Body:  copyStmtList(nodes, node.Body),
			}
		case *ast.ChanType:
			nodes[node] = &ast.ChanType{
				Begin: node.Begin,
				Arrow: node.Arrow,
				Dir:   node.Dir,
				Value: exprFromMap(nodes, node.Value),
			}
		case *ast.CommClause:
			nodes[node] = &ast.CommClause{
				Case:  node.Case,
				Comm:  stmtFromMap(nodes, node.Comm),
				Colon: node.Colon,
				Body:  copyStmtList(nodes, node.Body),
			}
		case *ast.Comment:
			nodes[node] = &ast.Comment{
				Slash: node.Slash,
				Text:  node.Text,
			}
		case *ast.CommentGroup:
			group := new(ast.CommentGroup)

			if node.List != nil {
				group.List = make([]*ast.Comment, len(node.List))

				for i := range node.List {
					group.List[i] = nodes[node.List[i]].(*ast.Comment)
				}
			}

			nodes[node] = group
		case *ast.CompositeLit:
			nodes[node] = &ast.CompositeLit{
				Type:   exprFromMap(nodes, node.Type),
				Lbrace: node.Lbrace,
				Elts:   copyExprList(nodes, node.Elts),
				Rbrace: node.Rbrace,
			}
		case *ast.DeclStmt:
			nodes[node] = &ast.DeclStmt{
				Decl: nodes[node.Decl].(ast.Decl),
			}
		case *ast.DeferStmt:
			nodes[node] = &ast.DeferStmt{
				Defer: node.Defer,
				Call:  callExprFromMap(nodes, node.Call),
			}
		case *ast.Ellipsis:
			nodes[node] = &ast.Ellipsis{
				Ellipsis: node.Ellipsis,
				Elt:      exprFromMap(nodes, node.Elt),
			}
		case *ast.EmptyStmt:
			nodes[node] = &ast.EmptyStmt{
				Semicolon: node.Semicolon,
				Implicit:  node.Implicit,
			}
		case *ast.ExprStmt:
			nodes[node] = &ast.ExprStmt{
				X: exprFromMap(nodes, node.X),
			}
		case *ast.Field:
			nodes[node] = &ast.Field{
				Doc:     commentGroupFromMap(nodes, node.Doc),
				Names:   copyIdentList(nodes, node.Names),
				Type:    exprFromMap(nodes, node.Type),
				Tag:     basicLitFromMap(nodes, node.Tag),
				Comment: commentGroupFromMap(nodes, node.Comment),
			}
		case *ast.FieldList:
			fiel := &ast.FieldList{
				Opening: node.Opening,
				Closing: node.Closing,
			}

			if node.List != nil {
				fiel.List = make([]*ast.Field, len(node.List))

				for i := range node.List {
					fiel.List[i] = nodes[node.List[i]].(*ast.Field)
				}
			}

			nodes[node] = fiel
		case *ast.ForStmt:
			nodes[node] = &ast.ForStmt{
				For:  node.For,
				Init: stmtFromMap(nodes, node.Init),
				Cond: exprFromMap(nodes, node.Cond),
				Post: stmtFromMap(nodes, node.Post),
				Body: blockStmtFromMap(nodes, node.Body),
			}
		case *ast.FuncDecl:
			nodes[node] = &ast.FuncDecl{
				Doc:  commentGroupFromMap(nodes, node.Doc),
				Recv: fieldListFromMap(nodes, node.Recv),
				Name: identFromMap(nodes, node.Name),
				Type: funcTypeFromMap(nodes, node.Type),
				Body: blockStmtFromMap(nodes, node.Body),
			}
		case *ast.FuncLit:
			nodes[node] = &ast.FuncLit{
				Type: funcTypeFromMap(nodes, node.Type),
				Body: blockStmtFromMap(nodes, node.Body),
			}
		case *ast.FuncType:
			nodes[node] = &ast.FuncType{
				Func:    node.Func,
				Params:  fieldListFromMap(nodes, node.Params),
				Results: fieldListFromMap(nodes, node.Results),
			}
		case *ast.GenDecl:
			decl := &ast.GenDecl{
				Doc:    commentGroupFromMap(nodes, node.Doc),
				TokPos: node.TokPos,
				Tok:    node.Tok,
				Lparen: node.Lparen,
				Rparen: node.Rparen,
			}

			if node.Specs != nil {
				decl.Specs = make([]ast.Spec, len(node.Specs))

				for i := range node.Specs {
					decl.Specs[i] = nodes[node.Specs[i]].(ast.Spec)
				}
			}

			nodes[node] = decl
		case *ast.GoStmt:
			nodes[node] = &ast.GoStmt{
				Go:   node.Go,
				Call: callExprFromMap(nodes, node.Call),
			}
		case *ast.Ident:
			// Keep identifiers the same identity so they can be conveniently
			// used with the original *types.Info.
			nodes[node] = node
		case *ast.IfStmt:
			nodes[node] = &ast.IfStmt{
				If:   node.If,
				Init: stmtFromMap(nodes, node.Init),
				Cond: exprFromMap(nodes, node.Cond),
				Body: blockStmtFromMap(nodes, node.Body),
				Else: stmtFromMap(nodes, node.Else),
			}
		case *ast.ImportSpec:
			nodes[node] = &ast.ImportSpec{
				Doc:     commentGroupFromMap(nodes, node.Doc),
				Name:    identFromMap(nodes, node.Name),
				Path:    basicLitFromMap(nodes, node.Path),
				Comment: commentGroupFromMap(nodes, node.Comment),
				EndPos:  node.EndPos,
			}
		case *ast.IncDecStmt:
			nodes[node] = &ast.IncDecStmt{
				X:      exprFromMap(nodes, node.X),
				TokPos: node.TokPos,
				Tok:    node.Tok,
			}
		case *ast.IndexExpr:
			nodes[node] = &ast.IndexExpr{
				X:      exprFromMap(nodes, node.X),
				Lbrack: node.Lbrack,
				Index:  exprFromMap(nodes, node.Index),
				Rbrack: node.Rbrack,
			}
		case *ast.InterfaceType:
			nodes[node] = &ast.InterfaceType{
				Interface:  node.Interface,
				Methods:    fieldListFromMap(nodes, node.Methods),
				Incomplete: node.Incomplete,
			}
		case *ast.KeyValueExpr:
			nodes[node] = &ast.KeyValueExpr{
				Key:   exprFromMap(nodes, node.Key),
				Colon: node.Colon,
				Value: exprFromMap(nodes, node.Value),
			}
		case *ast.LabeledStmt:
			nodes[node] = &ast.LabeledStmt{
				Label: identFromMap(nodes, node.Label),
				Colon: node.Colon,
				Stmt:  stmtFromMap(nodes, node.Stmt),
			}
		case *ast.MapType:
			nodes[node] = &ast.MapType{
				Map:   node.Map,
				Key:   exprFromMap(nodes, node.Key),
				Value: exprFromMap(nodes, node.Value),
			}
		case *ast.ParenExpr:
			nodes[node] = &ast.ParenExpr{
				Lparen: node.Lparen,
				X:      exprFromMap(nodes, node.X),
				Rparen: node.Rparen,
			}
		case *ast.RangeStmt:
			nodes[node] = &ast.RangeStmt{
				For:    node.For,
				Key:    exprFromMap(nodes, node.Key),
				Value:  exprFromMap(nodes, node.Value),
				TokPos: node.TokPos,
				Tok:    node.Tok,
				X:      exprFromMap(nodes, node.X),
				Body:   blockStmtFromMap(nodes, node.Body),
			}
		case *ast.ReturnStmt:
			nodes[node] = &ast.ReturnStmt{
				Return:  node.Return,
				Results: copyExprList(nodes, node.Results),
			}
		case *ast.SelectStmt:
			nodes[node] = &ast.SelectStmt{
				Select: node.Select,
				Body:   blockStmtFromMap(nodes, node.Body),
			}
		case *ast.SelectorExpr:
			nodes[node] = &ast.SelectorExpr{
				X:   exprFromMap(nodes, node.X),
				Sel: identFromMap(nodes, node.Sel),
			}
		case *ast.SendStmt:
			nodes[node] = &ast.SendStmt{
				Chan:  exprFromMap(nodes, node.Chan),
				Arrow: node.Arrow,
				Value: exprFromMap(nodes, node.Value),
			}
		case *ast.SliceExpr:
			nodes[node] = &ast.SliceExpr{
				X:      exprFromMap(nodes, node.X),
				Lbrack: node.Lbrack,
				Low:    exprFromMap(nodes, node.Low),
				High:   exprFromMap(nodes, node.High),
				Max:    exprFromMap(nodes, node.Max),
				Slice3: node.Slice3,
				Rbrack: node.Rbrack,
			}
		case *ast.StarExpr:
			nodes[node] = &ast.StarExpr{
				Star: node.Star,
				X:    exprFromMap(nodes, node.X),
			}
		case *ast.StructType:
			nodes[node] = &ast.StructType{
				Struct:     node.Struct,
				Fields:     fieldListFromMap(nodes, node.Fields),
				Incomplete: node.Incomplete,
			}
		case *ast.SwitchStmt:
			nodes[node] = &ast.SwitchStmt{
				Switch: node.Switch,
				Init:   stmtFromMap(nodes, node.Init),
				Tag:    exprFromMap(nodes, node.Tag),
				Body:   blockStmtFromMap(nodes, node.Body),
			}
		case *ast.TypeAssertExpr:
			nodes[node] = &ast.TypeAssertExpr{
				X:      exprFromMap(nodes, node.X),
				Lparen: node.Lparen,
				Type:   exprFromMap(nodes, node.Type),
				Rparen: node.Rparen,
			}
		case *ast.TypeSpec:
			nodes[node] = &ast.TypeSpec{
				Doc:     commentGroupFromMap(nodes, node.Doc),
				Name:    identFromMap(nodes, node.Name),
				Assign:  node.Assign,
				Type:    exprFromMap(nodes, node.Type),
				Comment: commentGroupFromMap(nodes, node.Comment),
			}
		case *ast.TypeSwitchStmt:
			nodes[node] = &ast.TypeSwitchStmt{
				Switch: node.Switch,
				Init:   stmtFromMap(nodes, node.Init),
				Assign: stmtFromMap(nodes, node.Assign),
				Body:   blockStmtFromMap(nodes, node.Body),
			}
		case *ast.UnaryExpr:
			nodes[node] = &ast.UnaryExpr{
				OpPos: node.OpPos,
				Op:    node.Op,
				X:     exprFromMap(nodes, node.X),
			}
		case *ast.ValueSpec:
			nodes[node] = &ast.ValueSpec{
				Doc:     commentGroupFromMap(nodes, node.Doc),
				Names:   copyIdentList(nodes, node.Names),
				Type:    exprFromMap(nodes, node.Type),
				Values:  copyExprList(nodes, node.Values),
				Comment: commentGroupFromMap(nodes, node.Comment),
			}
		default:
			panic(fmt.Sprintf("unhandled AST node: %T", node))
		}

		return true
	})

	return nodes[original]
}

func commentGroupFromMap(nodes map[ast.Node]ast.Node, key *ast.CommentGroup) *ast.CommentGroup {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.CommentGroup)
}

func exprFromMap(nodes map[ast.Node]ast.Node, key ast.Expr) ast.Expr {
	if key == nil {
		return nil
	}

	return nodes[key].(ast.Expr)
}

func stmtFromMap(nodes map[ast.Node]ast.Node, key ast.Stmt) ast.Stmt {
	if key == nil {
		return nil
	}

	return nodes[key].(ast.Stmt)
}

func identFromMap(nodes map[ast.Node]ast.Node, key *ast.Ident) *ast.Ident {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.Ident)
}

func blockStmtFromMap(nodes map[ast.Node]ast.Node, key *ast.BlockStmt) *ast.BlockStmt {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.BlockStmt)
}

func fieldListFromMap(nodes map[ast.Node]ast.Node, key *ast.FieldList) *ast.FieldList {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.FieldList)
}

func callExprFromMap(nodes map[ast.Node]ast.Node, key *ast.CallExpr) *ast.CallExpr {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.CallExpr)
}

func basicLitFromMap(nodes map[ast.Node]ast.Node, key *ast.BasicLit) *ast.BasicLit {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.BasicLit)
}

func funcTypeFromMap(nodes map[ast.Node]ast.Node, key *ast.FuncType) *ast.FuncType {
	if key == nil {
		return nil
	}

	return nodes[key].(*ast.FuncType)
}

func copyExprList(nodes map[ast.Node]ast.Node, exprs []ast.Expr) []ast.Expr {
	if exprs == nil {
		return nil
	}

	newExprs := make([]ast.Expr, len(exprs))
	for i := range exprs {
		newExprs[i] = nodes[exprs[i]].(ast.Expr)
	}

	return newExprs
}

func copyStmtList(nodes map[ast.Node]ast.Node, stmts []ast.Stmt) []ast.Stmt {
	if stmts == nil {
		return nil
	}

	newStmts := make([]ast.Stmt, len(stmts))
	for i := range stmts {
		newStmts[i] = nodes[stmts[i]].(ast.Stmt)
	}

	return newStmts
}

func copyIdentList(nodes map[ast.Node]ast.Node, idents []*ast.Ident) []*ast.Ident {
	if idents == nil {
		return nil
	}

	newIdents := make([]*ast.Ident, len(idents))
	for i := range idents {
		newIdents[i] = nodes[idents[i]].(*ast.Ident)
	}

	return newIdents
}