package loader

import (
	"fmt"
	"path/filepath"
	"strings"

	"github.com/marzeq/qk/attributes"
	"github.com/marzeq/qk/parser"
)

type PartialModuleInfo struct {
	Path                   string
	Name                   string
	Imports                []string
	Root                   *parser.RootNode
	Links                  []attributes.Link
	TrustedStandardLibrary bool
}

func CollectModuleInfo(root *parser.RootNode, trustedStandardLibrary bool) (*PartialModuleInfo, error) {
	name := ""
	imports := []string{}
	seenModule := false
	links := []attributes.Link{}

	for i, node := range root.Body {
		switch n := node.(type) {

		case *parser.ModuleNode:
			if i != 0 {
				return nil, fmt.Errorf("module declaration must be first statement")
			}
			if seenModule {
				return nil, fmt.Errorf("multiple module declarations")
			}
			seenModule = true
			name = n.Name
			for _, attr := range n.Attributes {
				linkAttr, ok := attr.(attributes.ModuleAttributeLink)
				if !ok {
					continue
				}
				for _, link := range linkAttr.Links {
					if (link.Kind == attributes.LinkPath || link.Kind == attributes.LinkSearchPath) && !filepath.IsAbs(link.Value) {
						link.Value = filepath.Join(filepath.Dir(n.Loc.FilePath), link.Value)
					}
					if link.Kind == attributes.LinkPath || link.Kind == attributes.LinkSearchPath {
						link.Value = filepath.Clean(link.Value)
					}
					links = append(links, link)
				}
			}

		case *parser.ImportNode:
			for i, imported := range n.Modules {
				if i < len(n.ResolvedModules) && n.ResolvedModules[i] != "" {
					imported = n.ResolvedModules[i]
				}
				imports = append(imports, imported)
			}
		}
	}

	if !seenModule || name == "" {
		return nil, fmt.Errorf("module declaration is missing or empty")
	}
	if (name == "std" || strings.HasPrefix(name, "std.")) && !trustedStandardLibrary {
		return nil, fmt.Errorf("module name %q is reserved for compiler-trusted standard-library sources", name)
	}
	if trustedStandardLibrary && name != "std" && !strings.HasPrefix(name, "std.") {
		return nil, fmt.Errorf("trusted standard-library source declares module %q outside the reserved %q namespace", name, "std")
	}

	return &PartialModuleInfo{
		Name:                   name,
		Imports:                imports,
		Root:                   root,
		Links:                  links,
		TrustedStandardLibrary: trustedStandardLibrary,
	}, nil
}

type ModuleInfo struct {
	Path                   string
	Name                   string
	Imports                []string
	Root                   *parser.RootNode
	Links                  []attributes.Link
	TrustedStandardLibrary bool
}

func BuildModules(partials []*PartialModuleInfo) (map[string]*ModuleInfo, error) {
	modules := map[string]*ModuleInfo{}

	for _, p := range partials {
		path := p.Path
		if path == "" {
			path = p.Name
		}
		if existing, ok := modules[path]; ok {
			if existing.Name != p.Name {
				return nil, fmt.Errorf("package %q contains both module %q and module %q", path, existing.Name, p.Name)
			}
			if existing.TrustedStandardLibrary != p.TrustedStandardLibrary {
				return nil, fmt.Errorf("cannot mix trusted and untrusted sources in package %q", path)
			}
			mergeModuleRoot(existing.Root, p.Root)
			existing.Imports = mergeImports(existing.Imports, p.Imports)
			existing.Links = append(existing.Links, p.Links...)
		} else {
			modules[path] = &ModuleInfo{
				Path:                   path,
				Name:                   p.Name,
				Imports:                unique(p.Imports),
				Root:                   cloneModuleRoot(p.Root),
				Links:                  append([]attributes.Link(nil), p.Links...),
				TrustedStandardLibrary: p.TrustedStandardLibrary,
			}
		}
	}

	return modules, nil
}

func cloneModuleRoot(root *parser.RootNode) *parser.RootNode {
	return &parser.RootNode{
		Body: append([]parser.Node(nil), root.Body...),
		Loc:  root.Loc,
	}
}

func mergeModuleRoot(destination, source *parser.RootNode) {
	for _, node := range source.Body {
		if _, isModule := node.(*parser.ModuleNode); isModule {
			continue
		}
		destination.Body = append(destination.Body, node)
	}
}

func mergeImports(a, b []string) []string {
	return unique(append(a, b...))
}

func unique(in []string) []string {
	seen := map[string]struct{}{}
	var out []string

	for _, v := range in {
		if _, ok := seen[v]; !ok {
			seen[v] = struct{}{}
			out = append(out, v)
		}
	}

	return out
}