// Copyright 2023 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package python

import (
	"context"
	_ "embed"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/emirpasic/gods/sets/treeset"
	godsutils "github.com/emirpasic/gods/utils"
	"golang.org/x/sync/errgroup"
)

// python3Parser implements a parser for Python files that extracts the modules
// as seen in the import statements.
type python3Parser struct {
	// The value of language.GenerateArgs.Config.RepoRoot.
	repoRoot string
	// The value of language.GenerateArgs.Rel.
	relPackagePath string
	// The function that determines if a dependency is ignored from a Gazelle
	// directive. It's the signature of pythonconfig.Config.IgnoresDependency.
	ignoresDependency func(dep string) bool
}

// newPython3Parser constructs a new python3Parser.
func newPython3Parser(
	repoRoot string,
	relPackagePath string,
	ignoresDependency func(dep string) bool,
) *python3Parser {
	return &python3Parser{
		repoRoot:          repoRoot,
		relPackagePath:    relPackagePath,
		ignoresDependency: ignoresDependency,
	}
}

// parseSingle parses a single Python file and returns the extracted modules
// from the import statements as well as the parsed comments.
func (p *python3Parser) parseSingle(pyFilename string) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {
	pyFilenames := treeset.NewWith(godsutils.StringComparator)
	pyFilenames.Add(pyFilename)
	return p.parse(pyFilenames)
}

// parse parses multiple Python files and returns the extracted modules from
// the import statements as well as the parsed comments.
func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {
	modules := treeset.NewWith(moduleComparator)

	g, ctx := errgroup.WithContext(context.Background())
	ch := make(chan struct{}, 6) // Limit the number of concurrent parses.
	chRes := make(chan *ParserOutput, len(pyFilenames.Values()))
	for _, v := range pyFilenames.Values() {
		ch <- struct{}{}
		g.Go(func(filename string) func() error {
			return func() error {
				defer func() {
					<-ch
				}()
				res, err := NewFileParser().ParseFile(ctx, p.repoRoot, p.relPackagePath, filename)
				if err != nil {
					return err
				}
				chRes <- res
				return nil
			}
		}(v.(string)))
	}
	if err := g.Wait(); err != nil {
		return nil, nil, nil, err
	}
	close(ch)
	close(chRes)
	mainModules := make(map[string]*treeset.Set, len(chRes))
	allAnnotations := new(annotations)
	allAnnotations.ignore = make(map[string]struct{})
	for res := range chRes {
		if res.HasMain {
			mainModules[res.FileName] = treeset.NewWith(moduleComparator)
		}
		annotations, err := annotationsFromComments(res.Comments)
		if err != nil {
			return nil, nil, nil, fmt.Errorf("failed to parse annotations: %w", err)
		}

		for _, m := range res.Modules {
			// Check for ignored dependencies set via an annotation to the Python
			// module.
			if annotations.ignores(m.Name) || annotations.ignores(m.From) {
				continue
			}

			// Check for ignored dependencies set via a Gazelle directive in a BUILD
			// file.
			if p.ignoresDependency(m.Name) || p.ignoresDependency(m.From) {
				continue
			}

			addModuleToTreeSet(modules, m)
			if res.HasMain {
				addModuleToTreeSet(mainModules[res.FileName], m)
			}
		}

		// Collect all annotations from each file into a single annotations struct.
		for k, v := range annotations.ignore {
			allAnnotations.ignore[k] = v
		}
		allAnnotations.includeDeps = append(allAnnotations.includeDeps, annotations.includeDeps...)
		allAnnotations.includePytestConftest = annotations.includePytestConftest
	}

	allAnnotations.includeDeps = removeDupesFromStringTreeSetSlice(allAnnotations.includeDeps)

	return modules, mainModules, allAnnotations, nil
}

// removeDupesFromStringTreeSetSlice takes a []string, makes a set out of the
// elements, and then returns a new []string with all duplicates removed. Order
// is preserved.
func removeDupesFromStringTreeSetSlice(array []string) []string {
	s := treeset.NewWith(godsutils.StringComparator)
	for _, v := range array {
		s.Add(v)
	}
	dedupe := make([]string, s.Size())
	for i, v := range s.Values() {
		dedupe[i] = fmt.Sprint(v)
	}
	return dedupe
}

// Module represents a fully-qualified, dot-separated, Python module as seen on
// the import statement, alongside the line number where it happened.
type Module struct {
	// The fully-qualified, dot-separated, Python module name as seen on import
	// statements.
	Name string `json:"name"`
	// The line number where the import happened.
	LineNumber uint32 `json:"lineno"`
	// The path to the module file relative to the Bazel workspace root.
	Filepath string `json:"filepath"`
	// If this was a from import, e.g. from foo import bar, From indicates the module
	// from which it is imported.
	From string `json:"from"`
	// Whether this import is type-checking only (inside if TYPE_CHECKING block).
	TypeCheckingOnly bool `json:"type_checking_only"`
}

// moduleComparator compares modules by name.
func moduleComparator(a, b interface{}) int {
	return godsutils.StringComparator(a.(Module).Name, b.(Module).Name)
}

// addModuleToTreeSet adds a module to a treeset.Set, ensuring that a TypeCheckingOnly=false module is
// prefered over a TypeCheckingOnly=true module.
func addModuleToTreeSet(set *treeset.Set, mod Module) {
	if mod.TypeCheckingOnly && set.Contains(mod) {
		return
	}
	set.Add(mod)
}

// annotationKind represents Gazelle annotation kinds.
type annotationKind string

const (
	// The Gazelle annotation prefix.
	annotationPrefix string = "gazelle:"
	// The ignore annotation kind. E.g. '# gazelle:ignore <module_name>'.
	annotationKindIgnore annotationKind = "ignore"
	// Force a particular target to be added to `deps`. Multiple invocations are
	// accumulated and the value can be comma separated.
	// Eg: '# gazelle:include_dep //foo/bar:baz,@repo//:target
	annotationKindIncludeDep            annotationKind = "include_dep"
	annotationKindIncludePytestConftest annotationKind = "include_pytest_conftest"
)

// Comment represents a Python comment.
type Comment string

// asAnnotation returns an annotation object if the comment has the
// annotationPrefix.
func (c *Comment) asAnnotation() (*annotation, error) {
	uncomment := strings.TrimLeft(string(*c), "# ")
	if !strings.HasPrefix(uncomment, annotationPrefix) {
		return nil, nil
	}
	withoutPrefix := strings.TrimPrefix(uncomment, annotationPrefix)
	annotationParts := strings.SplitN(withoutPrefix, " ", 2)
	if len(annotationParts) < 2 {
		return nil, fmt.Errorf("`%s` requires a value", *c)
	}
	return &annotation{
		kind:  annotationKind(annotationParts[0]),
		value: annotationParts[1],
	}, nil
}

// annotation represents a single Gazelle annotation parsed from a Python
// comment.
type annotation struct {
	kind  annotationKind
	value string
}

// annotations represent the collection of all Gazelle annotations parsed out of
// the comments of a Python module.
type annotations struct {
	// The parsed modules to be ignored by Gazelle.
	ignore map[string]struct{}
	// Labels that Gazelle should include as deps of the generated target.
	includeDeps []string
	// Whether the conftest.py file, found in the same directory as the current
	// python test file, should be added to the py_test target's `deps` attribute.
	// A *bool is used so that we can handle the "not set" state.
	includePytestConftest *bool
}

// annotationsFromComments returns all the annotations parsed out of the
// comments of a Python module.
func annotationsFromComments(comments []Comment) (*annotations, error) {
	ignore := make(map[string]struct{})
	includeDeps := []string{}
	var includePytestConftest *bool
	for _, comment := range comments {
		annotation, err := comment.asAnnotation()
		if err != nil {
			return nil, err
		}
		if annotation != nil {
			if annotation.kind == annotationKindIgnore {
				modules := strings.Split(annotation.value, ",")
				for _, m := range modules {
					if m == "" {
						continue
					}
					m = strings.TrimSpace(m)
					ignore[m] = struct{}{}
				}
			}
			if annotation.kind == annotationKindIncludeDep {
				targets := strings.Split(annotation.value, ",")
				for _, t := range targets {
					if t == "" {
						continue
					}
					t = strings.TrimSpace(t)
					includeDeps = append(includeDeps, t)
				}
			}
			if annotation.kind == annotationKindIncludePytestConftest {
				val := annotation.value
				parsedVal, err := strconv.ParseBool(val)
				if err != nil {
					log.Printf("WARNING: unable to cast %q to bool in %q. Ignoring annotation", val, comment)
					continue
				}
				includePytestConftest = &parsedVal
			}
		}
	}
	return &annotations{
		ignore:                ignore,
		includeDeps:           includeDeps,
		includePytestConftest: includePytestConftest,
	}, nil
}

// ignored returns true if the given module was ignored via the ignore
// annotation.
func (a *annotations) ignores(module string) bool {
	_, ignores := a.ignore[module]
	return ignores
}