/*
 * Copyright (C) 2024 The Android Open Source Project
 *
 * 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.
 */

// Binary ide_query generates and analyzes build artifacts.
// The produced result can be consumed by IDEs to provide language features.
package main

import (
	"bytes"
	"cmp"
	"container/list"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path"
	"slices"
	"strings"

	"google.golang.org/protobuf/proto"
	apb "ide_query/cc_analyzer_proto"
	pb "ide_query/ide_query_proto"
)

// Env contains information about the current environment.
type Env struct {
	LunchTarget    LunchTarget
	RepoDir        string
	OutDir         string
	ClangToolsRoot string
}

// LunchTarget is a parsed Android lunch target.
// Input format: <product_name>-<release_type>-<build_variant>
type LunchTarget struct {
	Product string
	Release string
	Variant string
}

var _ flag.Value = (*LunchTarget)(nil)

// // Get implements flag.Value.
// func (l *LunchTarget) Get() any {
// 	return l
// }

// Set implements flag.Value.
func (l *LunchTarget) Set(s string) error {
	parts := strings.Split(s, "-")
	if len(parts) != 3 {
		return fmt.Errorf("invalid lunch target: %q, must have form <product_name>-<release_type>-<build_variant>", s)
	}
	*l = LunchTarget{
		Product: parts[0],
		Release: parts[1],
		Variant: parts[2],
	}
	return nil
}

// String implements flag.Value.
func (l *LunchTarget) String() string {
	return fmt.Sprintf("%s-%s-%s", l.Product, l.Release, l.Variant)
}

func main() {
	var env Env
	env.OutDir = strings.TrimSuffix(os.Getenv("OUT_DIR"), "/")
	env.RepoDir = os.Getenv("ANDROID_BUILD_TOP")
	env.ClangToolsRoot = os.Getenv("PREBUILTS_CLANG_TOOLS_ROOT")
	flag.Var(&env.LunchTarget, "lunch_target", "The lunch target to query")
	flag.Parse()
	files := flag.Args()
	if len(files) == 0 {
		fmt.Println("No files provided.")
		os.Exit(1)
		return
	}

	var ccFiles, javaFiles []string
	for _, f := range files {
		switch {
		case strings.HasSuffix(f, ".java") || strings.HasSuffix(f, ".kt"):
			javaFiles = append(javaFiles, f)
		case strings.HasSuffix(f, ".cc") || strings.HasSuffix(f, ".cpp") || strings.HasSuffix(f, ".h"):
			ccFiles = append(ccFiles, f)
		default:
			log.Printf("File %q is supported - will be skipped.", f)
		}
	}

	ctx := context.Background()
	// TODO(michaelmerg): Figure out if module_bp_java_deps.json and compile_commands.json is outdated.
	runMake(ctx, env, "nothing")

	javaModules, err := loadJavaModules(env)
	if err != nil {
		log.Printf("Failed to load java modules: %v", err)
	}

	var targets []string
	javaTargetsByFile := findJavaModules(javaFiles, javaModules)
	for _, target := range javaTargetsByFile {
		targets = append(targets, javaModules[target].Jars...)
	}

	ccTargets, err := getCCTargets(ctx, env, ccFiles)
	if err != nil {
		log.Fatalf("Failed to query cc targets: %v", err)
	}
	targets = append(targets, ccTargets...)
	if len(targets) == 0 {
		fmt.Println("No targets found.")
		os.Exit(1)
		return
	}

	fmt.Fprintf(os.Stderr, "Running make for modules: %v\n", strings.Join(targets, ", "))
	if err := runMake(ctx, env, targets...); err != nil {
		log.Printf("Building modules failed: %v", err)
	}

	var analysis pb.IdeAnalysis
	results, units := getJavaInputs(env, javaTargetsByFile, javaModules)
	analysis.Results = results
	analysis.Units = units
	if err != nil && analysis.Error == nil {
		analysis.Error = &pb.AnalysisError{
			ErrorMessage: err.Error(),
		}
	}

	results, units, err = getCCInputs(ctx, env, ccFiles)
	analysis.Results = append(analysis.Results, results...)
	analysis.Units = append(analysis.Units, units...)
	if err != nil && analysis.Error == nil {
		analysis.Error = &pb.AnalysisError{
			ErrorMessage: err.Error(),
		}
	}

	analysis.BuildOutDir = env.OutDir
	data, err := proto.Marshal(&analysis)
	if err != nil {
		log.Fatalf("Failed to marshal result proto: %v", err)
	}

	_, err = os.Stdout.Write(data)
	if err != nil {
		log.Fatalf("Failed to write result proto: %v", err)
	}

	for _, r := range analysis.Results {
		fmt.Fprintf(os.Stderr, "%s: %+v\n", r.GetSourceFilePath(), r.GetStatus())
	}
}

func repoState(env Env, filePaths []string) *apb.RepoState {
	const compDbPath = "soong/development/ide/compdb/compile_commands.json"
	return &apb.RepoState{
		RepoDir:        env.RepoDir,
		ActiveFilePath: filePaths,
		OutDir:         env.OutDir,
		CompDbPath:     path.Join(env.OutDir, compDbPath),
	}
}

func runCCanalyzer(ctx context.Context, env Env, mode string, in []byte) ([]byte, error) {
	ccAnalyzerPath := path.Join(env.ClangToolsRoot, "bin/ide_query_cc_analyzer")
	outBuffer := new(bytes.Buffer)

	inBuffer := new(bytes.Buffer)
	inBuffer.Write(in)

	cmd := exec.CommandContext(ctx, ccAnalyzerPath, "--mode="+mode)
	cmd.Dir = env.RepoDir

	cmd.Stdin = inBuffer
	cmd.Stdout = outBuffer
	cmd.Stderr = os.Stderr

	err := cmd.Run()

	return outBuffer.Bytes(), err
}

// Execute cc_analyzer and get all the targets that needs to be build for analyzing files.
func getCCTargets(ctx context.Context, env Env, filePaths []string) ([]string, error) {
	state, err := proto.Marshal(repoState(env, filePaths))
	if err != nil {
		log.Fatalln("Failed to serialize state:", err)
	}

	resp := new(apb.DepsResponse)
	result, err := runCCanalyzer(ctx, env, "deps", state)
	if err != nil {
		return nil, err
	}

	if err := proto.Unmarshal(result, resp); err != nil {
		return nil, fmt.Errorf("malformed response from cc_analyzer: %v", err)
	}

	var targets []string
	if resp.Status != nil && resp.Status.Code != apb.Status_OK {
		return targets, fmt.Errorf("cc_analyzer failed: %v", resp.Status.Message)
	}

	for _, deps := range resp.Deps {
		targets = append(targets, deps.BuildTarget...)
	}
	return targets, nil
}

func getCCInputs(ctx context.Context, env Env, filePaths []string) ([]*pb.AnalysisResult, []*pb.BuildableUnit, error) {
	state, err := proto.Marshal(repoState(env, filePaths))
	if err != nil {
		log.Fatalln("Failed to serialize state:", err)
	}

	resp := new(apb.IdeAnalysis)
	result, err := runCCanalyzer(ctx, env, "inputs", state)
	if err != nil {
		return nil, nil, fmt.Errorf("cc_analyzer failed:", err)
	}
	if err := proto.Unmarshal(result, resp); err != nil {
		return nil, nil, fmt.Errorf("malformed response from cc_analyzer: %v", err)
	}
	if resp.Status != nil && resp.Status.Code != apb.Status_OK {
		return nil, nil, fmt.Errorf("cc_analyzer failed: %v", resp.Status.Message)
	}

	var results []*pb.AnalysisResult
	var units []*pb.BuildableUnit
	for _, s := range resp.Sources {
		status := &pb.AnalysisResult_Status{
			Code: pb.AnalysisResult_Status_CODE_OK,
		}
		if s.GetStatus().GetCode() != apb.Status_OK {
			status.Code = pb.AnalysisResult_Status_CODE_BUILD_FAILED
			status.StatusMessage = proto.String(s.GetStatus().GetMessage())
		}

		result := &pb.AnalysisResult{
			SourceFilePath: s.GetPath(),
			UnitId:         s.GetPath(),
			Status:         status,
		}
		results = append(results, result)

		var generated []*pb.GeneratedFile
		for _, f := range s.Generated {
			generated = append(generated, &pb.GeneratedFile{
				Path:     f.GetPath(),
				Contents: f.GetContents(),
			})
		}
		genUnit := &pb.BuildableUnit{
			Id:              "genfiles_for_" + s.GetPath(),
			SourceFilePaths: s.GetDeps(),
			GeneratedFiles:  generated,
		}

		unit := &pb.BuildableUnit{
			Id:                s.GetPath(),
			Language:          pb.Language_LANGUAGE_CPP,
			SourceFilePaths:   []string{s.GetPath()},
			CompilerArguments: s.GetCompilerArguments(),
			DependencyIds:     []string{genUnit.GetId()},
		}
		units = append(units, unit, genUnit)
	}
	return results, units, nil
}

// findJavaModules tries to find the modules that cover the given file paths.
// If a file is covered by multiple modules, the first module is returned.
func findJavaModules(paths []string, modules map[string]*javaModule) map[string]string {
	ret := make(map[string]string)
	// A file may be part of multiple modules. To make the result deterministic,
	// check the modules in sorted order.
	keys := make([]string, 0, len(modules))
	for name := range modules {
		keys = append(keys, name)
	}
	slices.SortFunc(keys, func(k1, k2 string) int {
		// Some libraries use annotations by setting annotations_enabled: true.
		// To enable IDE integration for such libraries, prefer the non package-private
		// version of the annotations library.
		if k1 == "stub-annotations" && k2 == "private-stub-annotations" {
			return -1
		} else if k1 == "private-stub-annotations" && k2 == "stub-annotations" {
			return 1
		} else {
			return cmp.Compare(k1, k2)
		}
	})
	for _, name := range keys {
		if strings.HasSuffix(name, ".impl") {
			continue
		}

		module := modules[name]
		if len(module.Jars) == 0 {
			continue
		}

		for i, p := range paths {
			if slices.Contains(module.Srcs, p) {
				ret[p] = name
				paths = append(paths[:i], paths[i+1:]...)
				break
			}
		}
		if len(paths) == 0 {
			break
		}
	}

	return ret
}

func getJavaInputs(env Env, modulesByPath map[string]string, modules map[string]*javaModule) ([]*pb.AnalysisResult, []*pb.BuildableUnit) {
	var results []*pb.AnalysisResult
	unitsById := make(map[string]*pb.BuildableUnit)
	for p, moduleName := range modulesByPath {
		r := &pb.AnalysisResult{
			SourceFilePath: p,
		}
		results = append(results, r)

		m := modules[moduleName]
		if m == nil {
			r.Status = &pb.AnalysisResult_Status{
				Code:          pb.AnalysisResult_Status_CODE_NOT_FOUND,
				StatusMessage: proto.String("File not found in any module."),
			}
			continue
		}

		r.UnitId = moduleName
		r.Status = &pb.AnalysisResult_Status{Code: pb.AnalysisResult_Status_CODE_OK}
		if unitsById[r.UnitId] != nil {
			// File is covered by an already created unit.
			continue
		}

		u := &pb.BuildableUnit{
			Id:              moduleName,
			Language:        pb.Language_LANGUAGE_JAVA,
			SourceFilePaths: m.Srcs,
			GeneratedFiles:  genFiles(env, m),
			DependencyIds:   m.Deps,
		}
		unitsById[u.Id] = u

		q := list.New()
		for _, d := range m.Deps {
			q.PushBack(d)
		}
		for q.Len() > 0 {
			name := q.Remove(q.Front()).(string)
			mod := modules[name]
			if mod == nil || unitsById[name] != nil {
				continue
			}

			unitsById[name] = &pb.BuildableUnit{
				Id:              name,
				SourceFilePaths: mod.Srcs,
				GeneratedFiles:  genFiles(env, mod),
				DependencyIds:   mod.Deps,
			}

			for _, d := range mod.Deps {
				q.PushBack(d)
			}
		}
	}

	units := make([]*pb.BuildableUnit, 0, len(unitsById))
	for _, u := range unitsById {
		units = append(units, u)
	}
	return results, units
}

// genFiles returns the generated files (paths that start with outDir/) for the
// given module. Generated files that do not exist are ignored.
func genFiles(env Env, mod *javaModule) []*pb.GeneratedFile {
	var paths []string
	paths = append(paths, mod.Srcs...)
	paths = append(paths, mod.SrcJars...)
	paths = append(paths, mod.Jars...)

	prefix := env.OutDir + "/"
	var ret []*pb.GeneratedFile
	for _, p := range paths {
		relPath, ok := strings.CutPrefix(p, prefix)
		if !ok {
			continue
		}

		contents, err := os.ReadFile(path.Join(env.RepoDir, p))
		if err != nil {
			continue
		}

		ret = append(ret, &pb.GeneratedFile{
			Path:     relPath,
			Contents: contents,
		})
	}
	return ret
}

// runMake runs Soong build for the given modules.
func runMake(ctx context.Context, env Env, modules ...string) error {
	args := []string{
		"--make-mode",
		"ANDROID_BUILD_ENVIRONMENT_CONFIG=googler-cog",
		"SOONG_GEN_COMPDB=1",
		"TARGET_PRODUCT=" + env.LunchTarget.Product,
		"TARGET_RELEASE=" + env.LunchTarget.Release,
		"TARGET_BUILD_VARIANT=" + env.LunchTarget.Variant,
		"TARGET_BUILD_TYPE=release",
		"-k",
	}
	args = append(args, modules...)
	cmd := exec.CommandContext(ctx, "build/soong/soong_ui.bash", args...)
	cmd.Dir = env.RepoDir
	cmd.Stdout = os.Stderr
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

type javaModule struct {
	Path    []string `json:"path,omitempty"`
	Deps    []string `json:"dependencies,omitempty"`
	Srcs    []string `json:"srcs,omitempty"`
	Jars    []string `json:"jars,omitempty"`
	SrcJars []string `json:"srcjars,omitempty"`
}

func loadJavaModules(env Env) (map[string]*javaModule, error) {
	javaDepsPath := path.Join(env.RepoDir, env.OutDir, "soong/module_bp_java_deps.json")
	data, err := os.ReadFile(javaDepsPath)
	if err != nil {
		return nil, err
	}

	var ret map[string]*javaModule // module name -> module
	if err = json.Unmarshal(data, &ret); err != nil {
		return nil, err
	}

	// Add top level java_sdk_library for .impl modules.
	for name, module := range ret {
		if striped := strings.TrimSuffix(name, ".impl"); striped != name {
			ret[striped] = module
		}
	}
	return ret, nil
}