package store

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"

	milvuslifecycle "github.com/vllm-project/semantic-router/src/semantic-router/pkg/milvus"
)

const (
	DefaultMilvusCollection       = "router_replay_records"
	DefaultMilvusConsistencyLevel = "Session"
	DefaultMilvusShardNum         = 2
)

// MilvusStore implements Storage using Milvus as the backend.
// Records are stored as entities in a Milvus collection.
type MilvusStore struct {
	client         client.Client
	collectionName string
	ttl            time.Duration
	asyncWrites    bool
	asyncChan      chan asyncOp
	done           chan struct{}
	wg             sync.WaitGroup
	pendingWrites  map[string]struct{}
	mu             sync.RWMutex
}

// NewMilvusStore creates a new Milvus storage backend.
func NewMilvusStore(cfg *MilvusConfig, ttlSeconds int, asyncWrites bool) (*MilvusStore, error) {
	if cfg == nil {
		return nil, fmt.Errorf("milvus config is required")
	}

	if cfg.Address == "" {
		cfg.Address = "localhost:19530"
	}

	collectionName := cfg.CollectionName
	if collectionName == "" {
		collectionName = DefaultMilvusCollection
	}

	milvusClient, err := milvuslifecycle.Connect(context.Background(), client.Config{
		Address:  cfg.Address,
		Username: cfg.Username,
		Password: cfg.Password,
	}, 10*time.Second)
	if err != nil {
		return nil, err
	}

	store := &MilvusStore{
		client:         milvusClient,
		collectionName: collectionName,
		ttl:            time.Duration(ttlSeconds) * time.Second,
		asyncWrites:    asyncWrites,
		done:           make(chan struct{}),
		pendingWrites:  make(map[string]struct{}),
	}

	if err := milvuslifecycle.Retry(
		context.Background(),
		3,
		2*time.Second,
		"ensure router replay Milvus collection",
		func(ctx context.Context) error {
			collCtx, collCancel := context.WithTimeout(ctx, 20*time.Second)
			defer collCancel()
			return store.createCollection(collCtx, cfg)
		},
	); err != nil {
		return nil, err
	}

	if asyncWrites {
		store.asyncChan = make(chan asyncOp, 100)
		go store.asyncWriter()
	}

	return store, nil
}

// createCollection creates the Milvus collection if it doesn't exist.
//
//nolint:gocognit
func (m *MilvusStore) createCollection(ctx context.Context, cfg *MilvusConfig) error {
	return milvuslifecycle.EnsureCollectionLoadedWithHooksRetry(
		ctx,
		m.client,
		m.collectionName,
		func(innerCtx context.Context) error {
			// Create schema
			schema := &entity.Schema{
				CollectionName: m.collectionName,
				Description:    "Router replay records",
				Fields: []*entity.Field{
					{
						Name:       "id",
						DataType:   entity.FieldTypeVarChar,
						PrimaryKey: true,
						AutoID:     false,
						TypeParams: map[string]string{
							"max_length": "255",
						},
					},
					{
						Name:     "timestamp",
						DataType: entity.FieldTypeInt64,
					},
					{
						Name:     "data",
						DataType: entity.FieldTypeVarChar,
						TypeParams: map[string]string{
							"max_length": "65535",
						},
					},
					{
						Name:     "vector",
						DataType: entity.FieldTypeFloatVector,
						TypeParams: map[string]string{
							"dim": "2",
						},
					},
				},
			}

			shardNum := cfg.ShardNum
			if shardNum <= 0 {
				shardNum = DefaultMilvusShardNum
			}
			if shardNum > 2147483647 {
				shardNum = DefaultMilvusShardNum
			}

			// Create collection
			//nolint:gosec // shardNum is validated to be within int32 range
			if err := m.client.CreateCollection(innerCtx, schema, int32(shardNum)); err != nil {
				return err
			}

			// Create vector index (required by Milvus even for dummy vectors)
			vectorIdx, err := entity.NewIndexAUTOINDEX(entity.L2)
			if err != nil {
				return fmt.Errorf("failed to build vector index: %w", err)
			}

			if err := m.client.CreateIndex(innerCtx, m.collectionName, "vector", vectorIdx, false); err != nil {
				return fmt.Errorf("failed to create vector index: %w", err)
			}

			// Create index on timestamp for efficient querying
			timeIdx := entity.NewGenericIndex(
				"timestamp_idx",
				entity.Sorted,
				map[string]string{},
			)

			if err := m.client.CreateIndex(innerCtx, m.collectionName, "timestamp", timeIdx, false); err != nil {
				return fmt.Errorf("failed to create index: %w", err)
			}

			return nil
		},
		func(innerCtx context.Context) error {
			// Ensure required vector index exists (Milvus needs one before insert)
			idxes, err := m.client.DescribeIndex(innerCtx, m.collectionName, "vector")
			if err != nil {
				return fmt.Errorf("failed to describe vector index: %w", err)
			}

			if len(idxes) == 0 {
				vectorIdx, err := entity.NewIndexAUTOINDEX(entity.L2)
				if err != nil {
					return fmt.Errorf("failed to build vector index: %w", err)
				}

				if err := m.client.CreateIndex(innerCtx, m.collectionName, "vector", vectorIdx, false); err != nil {
					return fmt.Errorf("failed to create vector index: %w", err)
				}
			}
			return nil
		},
		milvuslifecycle.CollectionRetryOptions{},
	)
}

// asyncWriter processes async write operations.
func (m *MilvusStore) asyncWriter() {
	for {
		select {
		case op := <-m.asyncChan:
			err := op.fn()
			if op.err != nil {
				op.err <- err
			}
			m.wg.Done()
		case <-m.done:
			return
		}
	}
}

// Add inserts a new record into Milvus.
func (m *MilvusStore) Add(ctx context.Context, record Record) (string, error) {
	if record.ID == "" {
		id, err := generateID()
		if err != nil {
			return "", err
		}
		record.ID = id
	}

	if record.Timestamp.IsZero() {
		record.Timestamp = time.Now().UTC()
	}

	// Marshal record to JSON
	data, err := json.Marshal(record)
	if err != nil {
		return "", fmt.Errorf("failed to marshal record: %w", err)
	}

	fn := func() error {
		// Create columns
		idColumn := entity.NewColumnVarChar("id", []string{record.ID})
		timestampColumn := entity.NewColumnInt64("timestamp", []int64{record.Timestamp.Unix()})
		dataColumn := entity.NewColumnVarChar("data", []string{string(data)})
		vectorColumn := entity.NewColumnFloatVector("vector", 2, [][]float32{{0.0, 0.0}})

		// Insert
		_, err := m.client.Insert(ctx, m.collectionName, "", idColumn, timestampColumn, dataColumn, vectorColumn)
		return err
	}

	if m.asyncWrites {
		m.mu.Lock()
		m.pendingWrites[record.ID] = struct{}{}
		m.mu.Unlock()
		m.wg.Add(1)
		errChan := make(chan error, 1)
		m.asyncChan <- asyncOp{fn: func() error {
			err := fn()
			m.mu.Lock()
			delete(m.pendingWrites, record.ID)
			m.mu.Unlock()
			if err == nil {
				// Flush to ensure data is persisted
				_ = m.client.Flush(ctx, m.collectionName, false)
			}
			return err
		}, err: errChan}
		return record.ID, nil
	}

	if err := fn(); err != nil {
		return "", fmt.Errorf("failed to insert record: %w", err)
	}

	// Flush to ensure data is persisted
	if err := m.client.Flush(ctx, m.collectionName, false); err != nil {
		return "", fmt.Errorf("failed to flush: %w", err)
	}

	return record.ID, nil
}

// Get retrieves a record by ID from Milvus.
func (m *MilvusStore) Get(ctx context.Context, id string) (Record, bool, error) {
	// If async writes enabled and record is pending, wait for it
	if m.asyncWrites {
		m.mu.RLock()
		_, pending := m.pendingWrites[id]
		m.mu.RUnlock()

		if pending {
			// Wait for pending writes to complete
			m.wg.Wait()
		}
	}

	// Query by ID
	expr := fmt.Sprintf("id == '%s'", id)
	result, err := m.client.Query(ctx, m.collectionName, nil, expr, []string{"id", "timestamp", "data"})
	if err != nil {
		return Record{}, false, fmt.Errorf("failed to query record: %w", err)
	}

	if result.Len() == 0 {
		return Record{}, false, nil
	}

	// Get data field
	dataColumn := result.GetColumn("data")
	dataVarChar, ok := dataColumn.(*entity.ColumnVarChar)
	if !ok {
		return Record{}, false, fmt.Errorf("unexpected column type for data")
	}

	if len(dataVarChar.Data()) == 0 {
		return Record{}, false, nil
	}

	// Unmarshal record
	var record Record
	if err := json.Unmarshal([]byte(dataVarChar.Data()[0]), &record); err != nil {
		return Record{}, false, fmt.Errorf("failed to unmarshal record: %w", err)
	}

	return record, true, nil
}

// List returns all records ordered by timestamp descending.
func (m *MilvusStore) List(ctx context.Context) ([]Record, error) {
	// If async writes enabled, wait for all pending writes to complete
	if m.asyncWrites {
		m.wg.Wait()
	}

	// Query all records with a high limit (empty expression requires limit in Milvus)
	result, err := m.client.Query(
		ctx,
		m.collectionName,
		nil,
		"",
		[]string{"id", "timestamp", "data"},
		client.WithLimit(10000),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to query records: %w", err)
	}

	if result.Len() == 0 {
		return []Record{}, nil
	}

	// Get data column
	dataColumn := result.GetColumn("data")
	dataVarChar, ok := dataColumn.(*entity.ColumnVarChar)
	if !ok {
		return nil, fmt.Errorf("unexpected column type for data")
	}

	// Parse all records
	records := make([]Record, 0, len(dataVarChar.Data()))
	for _, data := range dataVarChar.Data() {
		var record Record
		if err := json.Unmarshal([]byte(data), &record); err != nil {
			continue // Skip malformed records
		}
		records = append(records, record)
	}

	sort.Slice(records, func(i, j int) bool {
		return records[i].Timestamp.After(records[j].Timestamp)
	})

	return records, nil
}

// UpdateStatus updates the response status and flags for a record.
func (m *MilvusStore) UpdateStatus(ctx context.Context, id string, status int, fromCache bool, streaming bool) error {
	// Get existing record
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	// Update fields
	if status != 0 {
		record.ResponseStatus = status
	}
	record.FromCache = record.FromCache || fromCache
	record.Streaming = record.Streaming || streaming

	// Delete old record and insert updated one
	return m.upsertRecord(ctx, record)
}

// AttachRequest updates the request body for a record.
func (m *MilvusStore) AttachRequest(ctx context.Context, id string, body string, truncated bool) error {
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	record.RequestBody = body
	record.RequestBodyTruncated = record.RequestBodyTruncated || truncated

	return m.upsertRecord(ctx, record)
}

// AttachResponse updates the response body for a record.
func (m *MilvusStore) AttachResponse(ctx context.Context, id string, body string, truncated bool) error {
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	record.ResponseBody = body
	record.ResponseBodyTruncated = record.ResponseBodyTruncated || truncated

	return m.upsertRecord(ctx, record)
}

// AppendOutcome links post-route feedback to a record.
func (m *MilvusStore) AppendOutcome(ctx context.Context, id string, outcome Outcome) error {
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	record.Outcomes = append(record.Outcomes, cloneOutcome(outcome))
	return m.upsertRecord(ctx, record)
}

// UpdateHallucinationStatus updates hallucination detection results for a record.
func (m *MilvusStore) UpdateHallucinationStatus(ctx context.Context, id string, detected bool, confidence float32, spans []string) error {
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	record.HallucinationDetected = detected
	record.HallucinationConfidence = confidence
	record.HallucinationSpans = spans

	return m.upsertRecord(ctx, record)
}

// UpdateUsageCost updates token usage and pricing-derived cost fields for a record.
func (m *MilvusStore) UpdateUsageCost(ctx context.Context, id string, usage UsageCost) error {
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	record.PromptTokens = cloneIntPtr(usage.PromptTokens)
	record.CachedPromptTokens = cloneIntPtr(usage.CachedPromptTokens)
	record.CompletionTokens = cloneIntPtr(usage.CompletionTokens)
	record.TotalTokens = cloneIntPtr(usage.TotalTokens)
	record.ActualCost = cloneFloat64Ptr(usage.ActualCost)
	record.BaselineCost = cloneFloat64Ptr(usage.BaselineCost)
	record.CostSavings = cloneFloat64Ptr(usage.CostSavings)
	record.Currency = cloneStringPtr(usage.Currency)
	record.BaselineModel = cloneStringPtr(usage.BaselineModel)

	return m.upsertRecord(ctx, record)
}

// UpdateToolTrace updates tool-calling trace details for a record.
func (m *MilvusStore) UpdateToolTrace(ctx context.Context, id string, trace ToolTrace) error {
	record, found, err := m.Get(ctx, id)
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("record with ID %s not found", id)
	}

	record.ToolTrace = cloneToolTrace(&trace)
	return m.upsertRecord(ctx, record)
}

// upsertRecord updates a record by deleting and reinserting.
func (m *MilvusStore) upsertRecord(ctx context.Context, record Record) error {
	// Marshal record
	data, err := json.Marshal(record)
	if err != nil {
		return fmt.Errorf("failed to marshal record: %w", err)
	}

	fn := func() error {
		// Delete existing record
		expr := fmt.Sprintf("id == '%s'", record.ID)
		if err := m.client.Delete(ctx, m.collectionName, "", expr); err != nil {
			return fmt.Errorf("failed to delete old record: %w", err)
		}

		// Insert updated record
		idColumn := entity.NewColumnVarChar("id", []string{record.ID})
		timestampColumn := entity.NewColumnInt64("timestamp", []int64{record.Timestamp.Unix()})
		dataColumn := entity.NewColumnVarChar("data", []string{string(data)})
		vectorColumn := entity.NewColumnFloatVector("vector", 2, [][]float32{{0.0, 0.0}})

		_, err := m.client.Insert(ctx, m.collectionName, "", idColumn, timestampColumn, dataColumn, vectorColumn)
		if err != nil {
			return fmt.Errorf("failed to insert updated record: %w", err)
		}

		return m.client.Flush(ctx, m.collectionName, false)
	}

	if m.asyncWrites {
		m.wg.Add(1)
		m.asyncChan <- asyncOp{fn: func() error {
			err := fn()
			if err == nil {
				// Flush to ensure data is persisted
				_ = m.client.Flush(ctx, m.collectionName, false)
			}
			return err
		}}
		return nil
	}

	return fn()
}

// Close closes the Milvus client and stops async writer.
func (m *MilvusStore) Close() error {
	if m.asyncWrites {
		// Wait for all pending writes to complete
		m.wg.Wait()
		close(m.done)
	}
	return m.client.Close()
}
