package main

import (
        "bufio"
        "os"

        "shep/internal/fbs/shep"
        "shep/internal/protocol"
)

func main() {
        stdin := bufio.NewReader(os.Stdin)
        stdout := os.Stdout

        // Send handshake to host with declared capabilities
        capabilities := []protocol.CapabilityInfo{
                {Name: "inventory", Version: 1},
        }
        handshake := protocol.BuildHandshake("inventory", capabilities)
        if err := protocol.WriteMessage(stdout, protocol.KindHandshake, handshake); err != nil {
                os.Exit(1)
        }

        // Wait for handshake response
        kind, data, err := protocol.ReadMessage(stdin)
        if err != nil {
                os.Exit(1)
        }

        if kind != protocol.KindHandshakeResponse {
                os.Exit(1)
        }

        resp, err := protocol.ParseHandshakeResponse(data)
        if err != nil {
                os.Exit(1)
        }

        if resp.Status() != shep.StatusOk {
                os.Exit(1)
        }

        // Main message loop - handle capability invocations
        for {
                kind, data, err := protocol.ReadMessage(stdin)
                if err != nil {
                        break
                }

                switch kind {
                case protocol.KindEnvelope:
                        handleEnvelope(stdout, data)
                case protocol.KindShutdown:
                        return
                }
        }
}

func handleEnvelope(stdout *os.File, data []byte) {
        env, err := protocol.ParseEnvelope(data)
        if err != nil {
                return
        }

        capability := string(env.Capability())
        correlationID := string(env.CorrelationId())
        payload := env.PayloadBytes()

        // Route to capability handler
        var responsePayload []byte
        var responseErr string

        switch capability {
        case "inventory":
                responsePayload, responseErr = handleInventory(stdout, payload)
        default:
                responseErr = "unknown capability: " + capability
        }

        // Send response envelope
        response := protocol.BuildEnvelope(capability, correlationID, true, responseErr, responsePayload)
        protocol.WriteMessage(stdout, protocol.KindEnvelope, response)
}

func handleInventory(stdout *os.File, payload []byte) ([]byte, string) {
        // Parse the request
        req, err := ParseInventoryRequest(payload)
        if err != nil {
                return nil, "failed to parse request: " + err.Error()
        }

        // Log what we're doing
        logMsg := protocol.BuildLog(shep.LogLevelInfo, "parsing inventory: "+req.FilePath)
        protocol.WriteMessage(stdout, protocol.KindLog, logMsg)

        // Get the content to parse
        var content []byte
        if len(req.Content) > 0 {
                content = req.Content
        } else if req.FilePath != "" {
                content, err = os.ReadFile(req.FilePath)
                if err != nil {
                        return BuildInventoryErrorResponse("failed to read file: " + err.Error()), ""
                }
        } else {
                return BuildInventoryErrorResponse("no file path or content provided"), ""
        }

        // Parse the YAML inventory
        inv, warnings, err := ParseYAMLInventory(content)
        if err != nil {
                return BuildInventoryErrorResponse("failed to parse inventory: " + err.Error()), ""
        }

        // Build and return the response
        return BuildInventoryResponse(inv, warnings), ""
}

package main

import (
        "encoding/json"
        "fmt"
        "sort"

        flatbuffers "github.com/google/flatbuffers/go"
        "github.com/goccy/go-yaml"

        fbsdomain "shep/internal/fbs/shep/domain"
        fbsinventory "shep/internal/fbs/shep/cap/inventory"
)

// InventoryRequest is the parsed request from the host.
type InventoryRequest struct {
        FilePath string
        Content  []byte
}

// ParseInventoryRequest parses a FlatBuffer InventoryRequest.
func ParseInventoryRequest(data []byte) (*InventoryRequest, error) {
        if len(data) == 0 {
                return nil, fmt.Errorf("empty request data")
        }
        req := fbsinventory.GetRootAsInventoryRequest(data, 0)
        return &InventoryRequest{
                FilePath: string(req.FilePath()),
                Content:  req.ContentBytes(),
        }, nil
}

// Inventory represents the parsed inventory structure.
type Inventory struct {
        Hosts  []Host
        Groups []Group
}

// Host represents a host in the inventory.
type Host struct {
        Name string
        Vars map[string]string
}

// Group represents a group in the inventory.
type Group struct {
        Name     string
        Hosts    []string
        Children []string
        Vars     map[string]string
}

// ParseYAMLInventory parses an Ansible YAML inventory file.
func ParseYAMLInventory(content []byte) (*Inventory, []string, error) {
        var raw map[string]any
        if err := yaml.Unmarshal(content, &raw); err != nil {
                return nil, nil, fmt.Errorf("yaml unmarshal: %w", err)
        }

        inv := &Inventory{
                Hosts:  []Host{},
                Groups: []Group{},
        }
        var warnings []string

        // Track all hosts we've seen to avoid duplicates
        seenHosts := make(map[string]*Host)
        // Track all groups we've seen
        seenGroups := make(map[string]*Group)

        // Parse each top-level key as a group
        for groupName, groupData := range raw {
                groupWarnings := parseGroupRecursive(groupName, groupData, seenGroups, seenHosts)
                warnings = append(warnings, groupWarnings...)
        }

        // Convert maps to slices
        for _, h := range seenHosts {
                inv.Hosts = append(inv.Hosts, *h)
        }
        for _, g := range seenGroups {
                inv.Groups = append(inv.Groups, *g)
        }

        // Sort for deterministic output
        sort.Slice(inv.Hosts, func(i, j int) bool {
                return inv.Hosts[i].Name < inv.Hosts[j].Name
        })
        sort.Slice(inv.Groups, func(i, j int) bool {
                return inv.Groups[i].Name < inv.Groups[j].Name
        })

        return inv, warnings, nil
}

// parseGroupRecursive parses a group and all its nested children.
func parseGroupRecursive(name string, data any, seenGroups map[string]*Group, seenHosts map[string]*Host) []string {
        var warnings []string

        group := &Group{
                Name:     name,
                Hosts:    []string{},
                Children: []string{},
                Vars:     map[string]string{},
        }

        if data == nil {
                seenGroups[name] = group
                return warnings
        }

        groupMap, ok := data.(map[string]any)
        if !ok {
                warnings = append(warnings, fmt.Sprintf("group %q: expected map, got %T", name, data))
                seenGroups[name] = group
                return warnings
        }

        // Parse hosts
        if hostsData, ok := groupMap["hosts"]; ok {
                hostNames, hosts := parseHostsData(hostsData)
                group.Hosts = hostNames
                for _, h := range hosts {
                        if existing, ok := seenHosts[h.Name]; ok {
                                // Merge vars
                                for k, v := range h.Vars {
                                        existing.Vars[k] = v
                                }
                        } else {
                                hostCopy := h
                                seenHosts[h.Name] = &hostCopy
                        }
                }
        }

        // Parse children (nested groups)
        if childrenData, ok := groupMap["children"]; ok {
                childrenMap, ok := childrenData.(map[string]any)
                if ok {
                        for childName, childData := range childrenMap {
                                group.Children = append(group.Children, childName)
                                // Recursively parse child group
                                childWarnings := parseGroupRecursive(childName, childData, seenGroups, seenHosts)
                                warnings = append(warnings, childWarnings...)
                        }
                }
        }

        // Parse vars
        if varsData, ok := groupMap["vars"]; ok {
                group.Vars = parseVars(varsData)
        }

        // Sort for deterministic output
        sort.Strings(group.Hosts)
        sort.Strings(group.Children)

        seenGroups[name] = group
        return warnings
}

// parseHostsData parses the hosts section of a group.
func parseHostsData(data any) ([]string, []Host) {
        var hostNames []string
        var hosts []Host

        hostsMap, ok := data.(map[string]any)
        if !ok {
                return hostNames, hosts
        }

        for hostName, hostData := range hostsMap {
                hostNames = append(hostNames, hostName)

                host := Host{
                        Name: hostName,
                        Vars: map[string]string{},
                }

                // Host can have vars
                if hostData != nil {
                        if hostVars, ok := hostData.(map[string]any); ok {
                                host.Vars = parseVars(hostVars)
                        }
                }

                hosts = append(hosts, host)
        }

        sort.Strings(hostNames)
        return hostNames, hosts
}


// parseVars parses a vars map, converting values to strings.
func parseVars(data any) map[string]string {
        result := map[string]string{}

        varsMap, ok := data.(map[string]any)
        if !ok {
                return result
        }

        for k, v := range varsMap {
                result[k] = varToString(v)
        }

        return result
}

// varToString converts a variable value to string.
// Complex types (lists, maps) are JSON-encoded because FlatBuffers Var.value
// is a string field. This preserves structure while fitting the schema.
// TODO: When we implement typed Value union (string | int | bool | list | map),
// this function should return a proper typed value instead.
func varToString(v any) string {
        switch val := v.(type) {
        case string:
                return val
        case int, int64, float64:
                return fmt.Sprintf("%v", val)
        case bool:
                return fmt.Sprintf("%v", val)
        case nil:
                return ""
        default:
                // Complex types (lists, maps) are JSON-encoded
                b, err := json.Marshal(val)
                if err != nil {
                        return fmt.Sprintf("%v", val)
                }
                return string(b)
        }
}

// BuildInventoryResponse builds a FlatBuffer InventoryResponse.
func BuildInventoryResponse(inv *Inventory, warnings []string) []byte {
        builder := flatbuffers.NewBuilder(1024)

        // Build hosts
        hostOffsets := make([]flatbuffers.UOffsetT, len(inv.Hosts))
        for i, host := range inv.Hosts {
                hostOffsets[i] = buildHost(builder, &host)
        }

        fbsdomain.InventoryStartHostsVector(builder, len(hostOffsets))
        for i := len(hostOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(hostOffsets[i])
        }
        hostsVec := builder.EndVector(len(hostOffsets))

        // Build groups
        groupOffsets := make([]flatbuffers.UOffsetT, len(inv.Groups))
        for i, group := range inv.Groups {
                groupOffsets[i] = buildGroup(builder, &group)
        }

        fbsdomain.InventoryStartGroupsVector(builder, len(groupOffsets))
        for i := len(groupOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(groupOffsets[i])
        }
        groupsVec := builder.EndVector(len(groupOffsets))

        // Build Inventory
        fbsdomain.InventoryStart(builder)
        fbsdomain.InventoryAddHosts(builder, hostsVec)
        fbsdomain.InventoryAddGroups(builder, groupsVec)
        inventoryOffset := fbsdomain.InventoryEnd(builder)

        // Build warnings vector
        warningOffsets := make([]flatbuffers.UOffsetT, len(warnings))
        for i, w := range warnings {
                warningOffsets[i] = builder.CreateString(w)
        }

        fbsinventory.InventoryResponseStartWarningsVector(builder, len(warningOffsets))
        for i := len(warningOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(warningOffsets[i])
        }
        warningsVec := builder.EndVector(len(warningOffsets))

        // Build response
        fbsinventory.InventoryResponseStart(builder)
        fbsinventory.InventoryResponseAddInventory(builder, inventoryOffset)
        fbsinventory.InventoryResponseAddWarnings(builder, warningsVec)
        responseOffset := fbsinventory.InventoryResponseEnd(builder)

        builder.Finish(responseOffset)
        return builder.FinishedBytes()
}

// BuildInventoryErrorResponse builds a FlatBuffer InventoryResponse with an error.
func BuildInventoryErrorResponse(errMsg string) []byte {
        builder := flatbuffers.NewBuilder(256)

        errOffset := builder.CreateString(errMsg)

        fbsinventory.InventoryResponseStart(builder)
        fbsinventory.InventoryResponseAddError(builder, errOffset)
        responseOffset := fbsinventory.InventoryResponseEnd(builder)

        builder.Finish(responseOffset)
        return builder.FinishedBytes()
}

func buildHost(builder *flatbuffers.Builder, host *Host) flatbuffers.UOffsetT {
        nameOffset := builder.CreateString(host.Name)

        // Build vars
        varOffsets := make([]flatbuffers.UOffsetT, 0, len(host.Vars))
        for k, v := range host.Vars {
                varOffsets = append(varOffsets, buildVar(builder, k, v))
        }

        fbsdomain.HostStartVarsVector(builder, len(varOffsets))
        for i := len(varOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(varOffsets[i])
        }
        varsVec := builder.EndVector(len(varOffsets))

        fbsdomain.HostStart(builder)
        fbsdomain.HostAddName(builder, nameOffset)
        fbsdomain.HostAddVars(builder, varsVec)
        return fbsdomain.HostEnd(builder)
}

func buildGroup(builder *flatbuffers.Builder, group *Group) flatbuffers.UOffsetT {
        nameOffset := builder.CreateString(group.Name)

        // Build hosts vector
        hostOffsets := make([]flatbuffers.UOffsetT, len(group.Hosts))
        for i, h := range group.Hosts {
                hostOffsets[i] = builder.CreateString(h)
        }

        fbsdomain.GroupStartHostsVector(builder, len(hostOffsets))
        for i := len(hostOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(hostOffsets[i])
        }
        hostsVec := builder.EndVector(len(hostOffsets))

        // Build children vector
        childOffsets := make([]flatbuffers.UOffsetT, len(group.Children))
        for i, c := range group.Children {
                childOffsets[i] = builder.CreateString(c)
        }

        fbsdomain.GroupStartChildrenVector(builder, len(childOffsets))
        for i := len(childOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(childOffsets[i])
        }
        childrenVec := builder.EndVector(len(childOffsets))

        // Build vars
        varOffsets := make([]flatbuffers.UOffsetT, 0, len(group.Vars))
        for k, v := range group.Vars {
                varOffsets = append(varOffsets, buildVar(builder, k, v))
        }

        fbsdomain.GroupStartVarsVector(builder, len(varOffsets))
        for i := len(varOffsets) - 1; i >= 0; i-- {
                builder.PrependUOffsetT(varOffsets[i])
        }
        varsVec := builder.EndVector(len(varOffsets))

        fbsdomain.GroupStart(builder)
        fbsdomain.GroupAddName(builder, nameOffset)
        fbsdomain.GroupAddHosts(builder, hostsVec)
        fbsdomain.GroupAddChildren(builder, childrenVec)
        fbsdomain.GroupAddVars(builder, varsVec)
        return fbsdomain.GroupEnd(builder)
}

func buildVar(builder *flatbuffers.Builder, key, value string) flatbuffers.UOffsetT {
        keyOffset := builder.CreateString(key)
        valueOffset := builder.CreateString(value)

        fbsdomain.VarStart(builder)
        fbsdomain.VarAddKey(builder, keyOffset)
        fbsdomain.VarAddValue(builder, valueOffset)
        return fbsdomain.VarEnd(builder)
}