mapper/mapper.go

package mapper

import (
	"encoding/json"
	"fmt"

	"github.com/KorAP/KoralPipe-TermMapper/ast"
	"github.com/KorAP/KoralPipe-TermMapper/config"
	"github.com/KorAP/KoralPipe-TermMapper/matcher"
	"github.com/KorAP/KoralPipe-TermMapper/parser"
)

// Direction represents the mapping direction (A to B or B to A)
type Direction bool

const (
	AtoB Direction = true
	BtoA Direction = false
)

// String converts the Direction to its string representation
func (d Direction) String() string {
	if d {
		return "atob"
	}
	return "btoa"
}

// ParseDirection converts a string direction to Direction type
func ParseDirection(dir string) (Direction, error) {
	switch dir {
	case "atob":
		return AtoB, nil
	case "btoa":
		return BtoA, nil
	default:
		return false, fmt.Errorf("invalid direction: %s", dir)
	}
}

// Mapper handles the application of mapping rules to JSON objects
type Mapper struct {
	mappingLists map[string]*config.MappingList
	parsedRules  map[string][]*parser.MappingResult
}

// NewMapper creates a new Mapper instance from a list of MappingLists
func NewMapper(lists []config.MappingList) (*Mapper, error) {
	m := &Mapper{
		mappingLists: make(map[string]*config.MappingList),
		parsedRules:  make(map[string][]*parser.MappingResult),
	}

	// Store mapping lists by ID
	for _, list := range lists {
		if _, exists := m.mappingLists[list.ID]; exists {
			return nil, fmt.Errorf("duplicate mapping list ID found: %s", list.ID)
		}

		// Create a copy of the list to store
		listCopy := list
		m.mappingLists[list.ID] = &listCopy

		// Parse the rules immediately
		parsedRules, err := list.ParseMappings()
		if err != nil {
			return nil, fmt.Errorf("failed to parse mappings for list %s: %w", list.ID, err)
		}
		m.parsedRules[list.ID] = parsedRules
	}

	return m, nil
}

// MappingOptions contains the options for applying mappings
type MappingOptions struct {
	FoundryA    string
	LayerA      string
	FoundryB    string
	LayerB      string
	Direction   Direction
	AddRewrites bool
}

// ApplyQueryMappings applies the specified mapping rules to a JSON object
func (m *Mapper) ApplyQueryMappings(mappingID string, opts MappingOptions, jsonData any) (any, error) {
	// Validate mapping ID
	if _, exists := m.mappingLists[mappingID]; !exists {
		return nil, fmt.Errorf("mapping list with ID %s not found", mappingID)
	}

	// Get the parsed rules
	rules := m.parsedRules[mappingID]

	// Check if we have a wrapper object with a "query" field
	var queryData any
	var hasQueryWrapper bool

	if jsonMap, ok := jsonData.(map[string]any); ok {
		if query, exists := jsonMap["query"]; exists {
			queryData = query
			hasQueryWrapper = true
		}
	}

	// If no query wrapper was found, use the entire input
	if !hasQueryWrapper {
		// If the input itself is not a valid query object, return it as is
		if !isValidQueryObject(jsonData) {
			return jsonData, nil
		}
		queryData = jsonData
	} else if queryData == nil || !isValidQueryObject(queryData) {
		// If we have a query wrapper but the query is nil or not a valid object,
		// return the original data
		return jsonData, nil
	}

	// Store rewrites if they exist
	var oldRewrites any
	if queryMap, ok := queryData.(map[string]any); ok {
		if rewrites, exists := queryMap["rewrites"]; exists {
			oldRewrites = rewrites
			delete(queryMap, "rewrites")
		}
	}

	// Convert input JSON to AST
	jsonBytes, err := json.Marshal(queryData)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal input JSON: %w", err)
	}

	node, err := parser.ParseJSON(jsonBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse JSON into AST: %w", err)
	}

	// Store whether the input was a Token
	isToken := false
	var tokenWrap ast.Node
	if token, ok := node.(*ast.Token); ok {
		isToken = true
		tokenWrap = token.Wrap
		node = tokenWrap
	}

	// Store original node for rewrite if needed
	var originalNode ast.Node
	if opts.AddRewrites {
		originalBytes, err := parser.SerializeToJSON(node)
		if err != nil {
			return nil, fmt.Errorf("failed to serialize original node for rewrite: %w", err)
		}
		originalNode, err = parser.ParseJSON(originalBytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse original node for rewrite: %w", err)
		}
	}

	// Apply each rule to the AST
	for _, rule := range rules {
		// Create pattern and replacement based on direction
		var pattern, replacement ast.Node
		if opts.Direction { // true means AtoB
			pattern = rule.Upper
			replacement = rule.Lower
		} else {
			pattern = rule.Lower
			replacement = rule.Upper
		}

		// Extract the inner nodes from the pattern and replacement tokens
		if token, ok := pattern.(*ast.Token); ok {
			pattern = token.Wrap
		}
		if token, ok := replacement.(*ast.Token); ok {
			replacement = token.Wrap
		}

		// Apply foundry and layer overrides
		if opts.Direction { // true means AtoB
			applyFoundryAndLayerOverrides(pattern, opts.FoundryA, opts.LayerA)
			applyFoundryAndLayerOverrides(replacement, opts.FoundryB, opts.LayerB)
		} else {
			applyFoundryAndLayerOverrides(pattern, opts.FoundryB, opts.LayerB)
			applyFoundryAndLayerOverrides(replacement, opts.FoundryA, opts.LayerA)
		}

		// Create matcher and apply replacement
		m, err := matcher.NewMatcher(ast.Pattern{Root: pattern}, ast.Replacement{Root: replacement})
		if err != nil {
			return nil, fmt.Errorf("failed to create matcher: %w", err)
		}
		node = m.Replace(node)
	}

	// Wrap the result in a token if the input was a token
	var result ast.Node
	if isToken {
		result = &ast.Token{Wrap: node}
	} else {
		result = node
	}

	// Convert AST back to JSON
	resultBytes, err := parser.SerializeToJSON(result)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize AST to JSON: %w", err)
	}

	// Parse the JSON string back into
	var resultData any
	if err := json.Unmarshal(resultBytes, &resultData); err != nil {
		return nil, fmt.Errorf("failed to parse result JSON: %w", err)
	}

	// Add rewrites if enabled and node was changed
	if opts.AddRewrites && !ast.NodesEqual(node, originalNode) {
		// Create rewrite object
		rewrite := map[string]any{
			"@type":  "koral:rewrite",
			"editor": "termMapper",
		}

		// Check if all terms in a group have their foundry changed
		if term, ok := originalNode.(*ast.Term); ok {
			if termGroup, ok := node.(*ast.TermGroup); ok {
				// Check if all terms in the group have a different foundry
				allFoundryChanged := true
				for _, op := range termGroup.Operands {
					if t, ok := op.(*ast.Term); ok {
						if t.Foundry == term.Foundry {
							allFoundryChanged = false
							break
						}
					}
				}
				if allFoundryChanged {
					rewrite["scope"] = "foundry"
					rewrite["src"] = term.Foundry
				} else {
					// Full node replacement
					originalBytes, err := parser.SerializeToJSON(originalNode)
					if err != nil {
						return nil, fmt.Errorf("failed to serialize original node for rewrite: %w", err)
					}
					var originalJSON any
					if err := json.Unmarshal(originalBytes, &originalJSON); err != nil {
						return nil, fmt.Errorf("failed to parse original node JSON for rewrite: %w", err)
					}
					rewrite["src"] = originalJSON
				}
			} else if newTerm, ok := node.(*ast.Term); ok {
				// Single term changes
				if term.Foundry != newTerm.Foundry {
					rewrite["scope"] = "foundry"
					rewrite["src"] = term.Foundry
				} else if term.Layer != newTerm.Layer {
					rewrite["scope"] = "layer"
					rewrite["src"] = term.Layer
				} else if term.Key != newTerm.Key {
					rewrite["scope"] = "key"
					rewrite["src"] = term.Key
				} else if term.Value != newTerm.Value {
					rewrite["scope"] = "value"
					rewrite["src"] = term.Value
				} else {
					// No specific attribute changed, use full node replacement
					originalBytes, err := parser.SerializeToJSON(originalNode)
					if err != nil {
						return nil, fmt.Errorf("failed to serialize original node for rewrite: %w", err)
					}
					var originalJSON any
					if err := json.Unmarshal(originalBytes, &originalJSON); err != nil {
						return nil, fmt.Errorf("failed to parse original node JSON for rewrite: %w", err)
					}
					rewrite["src"] = originalJSON
				}
			}
		} else {
			// Full node replacement
			originalBytes, err := parser.SerializeToJSON(originalNode)
			if err != nil {
				return nil, fmt.Errorf("failed to serialize original node for rewrite: %w", err)
			}
			var originalJSON any
			if err := json.Unmarshal(originalBytes, &originalJSON); err != nil {
				return nil, fmt.Errorf("failed to parse original node JSON for rewrite: %w", err)
			}
			rewrite["src"] = originalJSON
		}

		// Add rewrite to the node
		if resultMap, ok := resultData.(map[string]any); ok {
			if wrapMap, ok := resultMap["wrap"].(map[string]any); ok {
				rewrites, exists := wrapMap["rewrites"]
				if !exists {
					rewrites = []any{}
				}
				if rewritesList, ok := rewrites.([]any); ok {
					wrapMap["rewrites"] = append(rewritesList, rewrite)
				} else {
					wrapMap["rewrites"] = []any{rewrite}
				}
			}
		}
	}

	// Restore rewrites if they existed
	if oldRewrites != nil {
		if resultMap, ok := resultData.(map[string]any); ok {
			resultMap["rewrites"] = oldRewrites
		}
	}

	// If we had a query wrapper, put the transformed data back in it
	if hasQueryWrapper {
		if wrapper, ok := jsonData.(map[string]any); ok {
			wrapper["query"] = resultData
			return wrapper, nil
		}
	}

	return resultData, nil
}

// isValidQueryObject checks if the query data is a valid object that can be processed
func isValidQueryObject(data any) bool {
	// Check if it's a map
	queryMap, ok := data.(map[string]any)
	if !ok {
		return false
	}

	// Check if it has the required @type field
	if _, ok := queryMap["@type"]; !ok {
		return false
	}

	return true
}

// applyFoundryAndLayerOverrides recursively applies foundry and layer overrides to terms
func applyFoundryAndLayerOverrides(node ast.Node, foundry, layer string) {
	if node == nil {
		return
	}

	switch n := node.(type) {
	case *ast.Term:
		if foundry != "" {
			n.Foundry = foundry
		}
		if layer != "" {
			n.Layer = layer
		}
	case *ast.TermGroup:
		for _, op := range n.Operands {
			applyFoundryAndLayerOverrides(op, foundry, layer)
		}
	case *ast.Token:
		if n.Wrap != nil {
			applyFoundryAndLayerOverrides(n.Wrap, foundry, layer)
		}
	case *ast.CatchallNode:
		if n.Wrap != nil {
			applyFoundryAndLayerOverrides(n.Wrap, foundry, layer)
		}
		for _, op := range n.Operands {
			applyFoundryAndLayerOverrides(op, foundry, layer)
		}
	}
}