package model

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

	ort "github.com/yalue/onnxruntime_go"
)

const defaultFeatureCount = 14

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

type Predictor struct {
	session      *ort.DynamicAdvancedSession
	threshold    float64
	featureCount int
	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
	}

	featureCount := len(meta.FeatureColumns)
	if featureCount == 0 {
		featureCount = defaultFeatureCount
	}

	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,
		featureCount: featureCount,
	}, nil
}

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

func (p *Predictor) FeatureCount() int {
	return p.featureCount
}

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

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

	n := len(features)
	flat := make([]float32, 0, n*p.featureCount)
	for i, row := range features {
		if len(row) != p.featureCount {
			return nil, fmt.Errorf(
				"fila %d: se esperaban %d features, se recibieron %d (revisa pipeline vs inference.json)",
				i, p.featureCount, len(row),
			)
		}
		flat = append(flat, row...)
	}

	inputShape := ort.NewShape(int64(n), int64(p.featureCount))
	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
}