package main

import (
	"database/sql"
	"log"
	"math"
	"sync"
)

const bucketDurationMs = 5000 // 5 seconds

// Aggregator accumulates ticks into 5-second buckets and writes completed
// feature rows to the features database.
type Aggregator struct {
	mu            sync.Mutex
	currentBucket int64     // epoch ms of current bucket start
	ticks         []Tick    // ticks within the current bucket
	featDB        *sql.DB   // connection to features.db
	lastPrice     float64   // closing price of the previous bucket (for log return)
}

// NewAggregator creates an aggregator with a connection to the features database.
func NewAggregator(sm *StorageManager) (*Aggregator, error) {
	db, err := sm.OpenFeaturesDB()
	if err != nil {
		return nil, err
	}

	// Attempt to load the last close price from features.db for continuity
	var lastPrice float64
	err = db.QueryRow(`
		SELECT close_price FROM five_second_features 
		ORDER BY timestamp DESC LIMIT 1
	`).Scan(&lastPrice)
	if err != nil {
		lastPrice = 0 // Will be set from first tick
	}

	return &Aggregator{
		featDB:    db,
		lastPrice: lastPrice,
	}, nil
}

// ProcessTick adds a tick to the current bucket. If the tick crosses a 5-second
// boundary, the previous bucket is aggregated and flushed to the database.
func (a *Aggregator) ProcessTick(tick Tick) {
	a.mu.Lock()
	defer a.mu.Unlock()

	// Determine which 5-second bucket this tick belongs to
	tickBucket := (tick.Timestamp / bucketDurationMs) * bucketDurationMs

	if a.currentBucket == 0 {
		// First tick ever — initialize the bucket
		a.currentBucket = tickBucket
	}

	if tickBucket > a.currentBucket {
		// This tick belongs to a new bucket — flush the previous one
		if len(a.ticks) > 0 {
			a.flushBucket()
		}
		a.currentBucket = tickBucket
		a.ticks = a.ticks[:0] // reset slice, keep backing array
	}

	a.ticks = append(a.ticks, tick)
}

// Close shuts down the aggregator, flushing any remaining bucket and closing the DB.
func (a *Aggregator) Close() {
	a.mu.Lock()
	defer a.mu.Unlock()

	if len(a.ticks) > 0 {
		a.flushBucket()
	}
	if a.featDB != nil {
		a.featDB.Close()
	}
}

// flushBucket computes features from the accumulated ticks and writes to features.db.
// Must be called while holding a.mu.
func (a *Aggregator) flushBucket() {
	if len(a.ticks) == 0 {
		return
	}

	bucket := a.computeFeatures(a.ticks)

	_, err := a.featDB.Exec(`
		INSERT OR IGNORE INTO five_second_features 
		(timestamp, log_return, realized_vol, ofi, volume_sum, close_price)
		VALUES (?, ?, ?, ?, ?, ?)
	`, bucket.Timestamp, bucket.LogReturn, bucket.RealizedVol,
		bucket.OFI, bucket.VolumeSum, bucket.ClosePrice)

	if err != nil {
		log.Printf("[aggregator] Failed to write feature bucket: %v", err)
	}

	// Update lastPrice for next bucket's log return calculation
	a.lastPrice = bucket.ClosePrice
}

// computeFeatures calculates all feature columns from a slice of ticks.
func (a *Aggregator) computeFeatures(ticks []Tick) FeatureBucket {
	n := len(ticks)
	startPrice := ticks[0].Price
	closePrice := ticks[n-1].Price

	// Log return: ln(close / start_of_bucket_or_prev_close)
	refPrice := a.lastPrice
	if refPrice == 0 {
		refPrice = startPrice
	}
	logReturn := math.Log(closePrice / refPrice)

	// Order Flow Imbalance (OFI) and total volume
	var buyVol, sellVol, volumeSum float64
	for _, t := range ticks {
		volumeSum += t.Volume
		if t.Side == "Buy" {
			buyVol += t.Volume
		} else {
			sellVol += t.Volume
		}
	}
	ofi := buyVol - sellVol

	// Realized volatility: standard deviation of tick-to-tick log returns
	realizedVol := 0.0
	if n > 1 {
		logReturns := make([]float64, n-1)
		for i := 1; i < n; i++ {
			if ticks[i-1].Price > 0 {
				logReturns[i-1] = math.Log(ticks[i].Price / ticks[i-1].Price)
			}
		}

		// Calculate mean
		var sum float64
		for _, lr := range logReturns {
			sum += lr
		}
		mean := sum / float64(len(logReturns))

		// Calculate variance
		var variance float64
		for _, lr := range logReturns {
			diff := lr - mean
			variance += diff * diff
		}
		variance /= float64(len(logReturns))
		realizedVol = math.Sqrt(variance)
	}

	return FeatureBucket{
		Timestamp:   a.currentBucket,
		LogReturn:   logReturn,
		RealizedVol: realizedVol,
		OFI:         ofi,
		VolumeSum:   volumeSum,
		ClosePrice:  closePrice,
	}
}