package model

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"sync"

	ort "github.com/yalue/onnxruntime_go"
)

type InferenceMeta struct {
	Threshold      float64  `json:"threshold"`
	FeatureColumns []string `json:"feature_columns"`
}

type Predictor struct {
	session   *ort.DynamicAdvancedSession
	threshold float64
	mu        sync.Mutex
}

func NewPredictor(modelPath, metaPath string) (*Predictor, error) {
	if lib := os.Getenv("ONNXRUNTIME_SHARED_LIBRARY"); lib != "" {
		ort.SetSharedLibraryPath(lib)
	}
	if err := ort.InitializeEnvironment(); err != nil {
		return nil, fmt.Errorf("onnx init: %w", err)
	}

	metaBytes, err := os.ReadFile(metaPath)
	if err != nil {
		return nil, fmt.Errorf("read inference meta: %w", err)
	}
	var meta InferenceMeta
	if err := json.Unmarshal(metaBytes, &meta); err != nil {
		return nil, fmt.Errorf("parse inference meta: %w", err)
	}
	if meta.Threshold == 0 {
		meta.Threshold = 0.5
	}

	session, err := ort.NewDynamicAdvancedSession(
		modelPath,
		[]string{"float_input"},
		[]string{"probabilities"},
		nil,
	)
	if err != nil {
		return nil, fmt.Errorf("onnx session: %w", err)
	}

	return &Predictor{session: session, threshold: meta.Threshold}, nil
}

func (p *Predictor) Threshold() float64 {
	return p.threshold
}

func (p *Predictor) Close() error {
	if p.session != nil {
		p.session.Destroy()
	}
	return ort.DestroyEnvironment()
}

func (p *Predictor) PredictProba(features [][8]float32) ([]float64, error) {
	if len(features) == 0 {
		return nil, nil
	}

	n := len(features)
	flat := make([]float32, 0, n*8)
	for _, row := range features {
		flat = append(flat, row[:]...)
	}

	inputShape := ort.NewShape(int64(n), 8)
	inputTensor, err := ort.NewTensor(inputShape, flat)
	if err != nil {
		return nil, err
	}
	defer inputTensor.Destroy()

	outputShape := ort.NewShape(int64(n), 2)
	outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
	if err != nil {
		return nil, err
	}
	defer outputTensor.Destroy()

	p.mu.Lock()
	err = p.session.Run([]ort.Value{inputTensor}, []ort.Value{outputTensor})
	p.mu.Unlock()
	if err != nil {
		return nil, err
	}

	data := outputTensor.GetData()
	probs := make([]float64, n)
	for i := 0; i < n; i++ {
		probs[i] = float64(data[i*2+1])
	}
	return probs, nil
}

func RoundProb(p float64) float64 {
	return math.Round(p*10000) / 10000
}