// issue-model trains and queries an explainable prototype model for normalized
// events. It uses deterministic pattern signatures as weak labels until human
// labels are available.
package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"
)

type event struct {
	EventID          string `json:"eventId"`
	Service          string `json:"service"`
	Kind             string `json:"kind"`
	Source           string `json:"source"`
	IsError          uint8  `json:"isError"`
	StatusClass      string `json:"statusClass,omitempty"`
	StatusCode       uint16 `json:"statusCode,omitempty"`
	Route            string `json:"route,omitempty"`
	SpanName         string `json:"spanName,omitempty"`
	MetricName       string `json:"metricName,omitempty"`
	MetricType       string `json:"metricType,omitempty"`
	ExceptionType    string `json:"exceptionType,omitempty"`
	Message          string `json:"messageNormalized,omitempty"`
	StackSignature   string `json:"stackSignature,omitempty"`
	SpanKind         uint8  `json:"spanKind,omitempty"`
	PeerIPSource     string `json:"peerIpSource,omitempty"`
	EdgePeerSource   string `json:"edgePeerSource,omitempty"`
	EdgeService      string `json:"edgeService,omitempty"`
	PatternSignature string `json:"patternSignature"`
}

type prototype struct {
	Pattern string             `json:"patternSignature"`
	Count   int                `json:"count"`
	Tokens  map[string]float64 `json:"tokens"`
}

type model struct {
	Version     string      `json:"version"`
	FeatureSpec string      `json:"featureSpec"`
	Prototypes  []prototype `json:"prototypes"`
}

type prediction struct {
	EventID       string   `json:"eventId"`
	ActualPattern string   `json:"actualPattern,omitempty"`
	Pattern       string   `json:"patternSignature"`
	Score         float64  `json:"score"`
	Matches       []string `json:"matches"`
	Rank          int      `json:"rank"`
	Correct       bool     `json:"correct,omitempty"`
	Accepted      bool     `json:"accepted"`
}

func main() {
	trainPath := flag.String("train", "", "JSONL training sample")
	modelPath := flag.String("model", "/tmp/issue-model-v001.json", "model path")
	predictPath := flag.String("predict", "", "JSONL events to score")
	top := flag.Int("top", 3, "number of predictions per event")
	minScore := flag.Float64("min-score", 0.75, "minimum top similarity to accept a match")
	minMargin := flag.Float64("min-margin", 0.05, "minimum gap between first and second match")
	flag.Parse()
	if *trainPath == "" && *predictPath == "" {
		log.Fatal("provide -train or -predict")
	}
	if *trainPath != "" {
		trained := train(*trainPath)
		writeJSON(*modelPath, trained)
		fmt.Printf("trained prototypes=%d model=%s\n", len(trained.Prototypes), *modelPath)
	}
	if *predictPath != "" {
		var trained model
		readJSON(*modelPath, &trained)
		predict(*predictPath, trained, *top, *minScore, *minMargin)
	}
}

func train(path string) model {
	groups := map[string]*prototype{}
	forEachEvent(path, func(item event) {
		if item.PatternSignature == "" {
			return
		}
		p := groups[item.PatternSignature]
		if p == nil {
			p = &prototype{Pattern: item.PatternSignature, Tokens: map[string]float64{}}
			groups[item.PatternSignature] = p
		}
		p.Count++
		for token := range features(item) {
			p.Tokens[token]++
		}
	})
	result := model{Version: "v1", FeatureSpec: "categorical tokens excluding IDs, timestamps, trace IDs, IPs, and raw values"}
	for _, p := range groups {
		for token, value := range p.Tokens {
			p.Tokens[token] = value / float64(p.Count)
		}
		result.Prototypes = append(result.Prototypes, *p)
	}
	sort.Slice(result.Prototypes, func(i, j int) bool { return result.Prototypes[i].Pattern < result.Prototypes[j].Pattern })
	return result
}

func predict(path string, trained model, top int, minScore, minMargin float64) {
	if top < 1 {
		top = 1
	}
	if len(trained.Prototypes) == 0 {
		log.Fatal("model contains no prototypes; train it with a non-empty normalized sample")
	}
	forEachEvent(path, func(item event) {
		input := features(item)
		results := make([]prediction, 0, len(trained.Prototypes))
		for _, p := range trained.Prototypes {
			score, matches := similarity(input, p.Tokens)
			results = append(results, prediction{EventID: item.EventID, Pattern: p.Pattern, Score: score, Matches: matches})
		}
		sort.Slice(results, func(i, j int) bool {
			if results[i].Score != results[j].Score {
				return results[i].Score > results[j].Score
			}
			return results[i].Pattern < results[j].Pattern
		})
		if len(results) > top {
			results = results[:top]
		}
		margin := results[0].Score
		if len(results) > 1 {
			margin -= results[1].Score
		}
		for i := range results {
			results[i].ActualPattern = item.PatternSignature
			results[i].Rank = i + 1
			results[i].Correct = results[i].Pattern == item.PatternSignature
			results[i].Accepted = results[0].Score >= minScore && margin >= minMargin
		}
		for _, result := range results {
			writeLine(result)
		}
	})
}

func features(item event) map[string]struct{} {
	result := map[string]struct{}{}
	add := func(key, value string) {
		if value != "" {
			result[key+"="+value] = struct{}{}
		}
	}
	add("service", item.Service)
	add("kind", item.Kind)
	add("source", item.Source)
	add("status", item.StatusClass)
	if item.StatusCode > 0 {
		add("statusCode", strconv.Itoa(int(item.StatusCode)))
	}
	add("route", item.Route)
	add("span", item.SpanName)
	add("metric", item.MetricName)
	add("metricType", item.MetricType)
	add("exception", item.ExceptionType)
	add("message", item.Message)
	add("stack", item.StackSignature)
	if item.SpanKind > 0 {
		add("spanKind", strconv.Itoa(int(item.SpanKind)))
	}
	add("peerSource", item.PeerIPSource)
	add("edgePeerSource", item.EdgePeerSource)
	add("edgeService", item.EdgeService)
	return result
}

func similarity(input map[string]struct{}, prototype map[string]float64) (float64, []string) {
	if len(input) == 0 || len(prototype) == 0 {
		return 0, nil
	}
	var dot, inputNorm, protoNorm float64
	matches := []string{}
	for token := range input {
		inputNorm++
		if weight, ok := prototype[token]; ok {
			dot += weight
			matches = append(matches, token)
		}
	}
	for _, weight := range prototype {
		protoNorm += weight * weight
	}
	if protoNorm == 0 {
		return 0, matches
	}
	score := dot / math.Sqrt(inputNorm*protoNorm)
	sort.Strings(matches)
	return score, matches
}

func forEachEvent(path string, callback func(event)) {
	file, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	scanner.Buffer(make([]byte, 64*1024), 4*1024*1024)
	for scanner.Scan() {
		var item event
		if err := json.Unmarshal(scanner.Bytes(), &item); err != nil {
			log.Fatal(err)
		}
		callback(item)
	}
	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}
}
func writeJSON(path string, value any) {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		log.Fatal(err)
	}
	if err := os.WriteFile(path, data, 0600); err != nil {
		log.Fatal(err)
	}
}
func readJSON(path string, value any) {
	data, err := os.ReadFile(path)
	if err != nil {
		log.Fatal(err)
	}
	if err := json.Unmarshal(data, value); err != nil {
		log.Fatal(err)
	}
}
func writeLine(value any) {
	data, err := json.Marshal(value)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(strings.TrimSpace(string(data)))
}