package volume

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/sandbox0-ai/sandbox0/infra-operator/api/config"
	obsmetrics "github.com/sandbox0-ai/sandbox0/pkg/observability/metrics"
	"github.com/sandbox0-ai/sandbox0/storage-proxy/pkg/db"
	"github.com/sandbox0-ai/sandbox0/storage-proxy/pkg/fsmeta"
	"github.com/sandbox0-ai/sandbox0/storage-proxy/pkg/s0fs"
	"github.com/sirupsen/logrus"
)

// MountRegistrar handles registration of volume mounts for distributed coordination
type MountRegistrar interface {
	RegisterMount(ctx context.Context, volumeID string, options MountOptions) error
	UnregisterMount(ctx context.Context, volumeID string) error
	ValidateMount(ctx context.Context, volumeID string, accessMode AccessMode) error
}

// VolumeContext holds the mounted runtime state for a volume.
type VolumeContext struct {
	VolumeID  string
	TeamID    string
	Backend   string
	S0FS      *s0fs.Engine
	Access    AccessMode
	MountedAt time.Time
	RootInode fsmeta.Ino
	RootPath  string
	CacheDir  string
	Observer  StorageObserver

	handleMu      sync.Mutex
	nextHandleID  uint64
	fileHandles   map[uint64]uint64
	dirHandleIDs  map[uint64]uint64
	openFileCount map[uint64]int
	unlinkedFiles map[uint64]struct{}

	materializeCancel context.CancelFunc
	materializeDone   chan struct{}
}

// MountSession tracks a single mount session on this instance.
type MountSession struct {
	ID        string
	TeamID    string
	CreatedAt time.Time
	Scope     MountSessionScope
}

type MountSessionScope string

const (
	MountSessionScopeUnknown MountSessionScope = ""
	MountSessionScopeDirect  MountSessionScope = "direct"
)

type invalidateTracker struct {
	pending map[string]struct{}
	done    chan struct{}
	err     error
}

type directMountLease struct {
	SessionID string
	InFlight  int
	LastUsed  time.Time
}

// Manager manages mounted volumes and mount sessions.
type Manager struct {
	mu             sync.RWMutex
	volumes        map[string]*VolumeContext
	mountSessions  map[string]map[string]*MountSession
	directMounts   map[string]*directMountLease
	invalidates    map[string]map[string]*invalidateTracker
	logger         *logrus.Logger
	config         *config.StorageProxyConfig
	metrics        *obsmetrics.StorageProxyMetrics
	backends       map[string]Backend
	defaultBackend string
	registrar      MountRegistrar // Optional: for distributed coordination
	observer       StorageObserver
}

// NewManager creates a new volume manager
func NewManager(logger *logrus.Logger, cfg *config.StorageProxyConfig, repo *db.Repository) *Manager {
	return NewManagerWithBackends(logger, cfg, map[string]Backend{
		BackendS0FS: NewS0FSBackend(logger, cfg, repo),
	}, DefaultBackendType())
}

// NewManagerWithBackend creates a manager with an explicit storage backend.
func NewManagerWithBackend(logger *logrus.Logger, cfg *config.StorageProxyConfig, repo *db.Repository, backend Backend) *Manager {
	if backend == nil {
		return NewManager(logger, cfg, repo)
	}
	return NewManagerWithBackends(logger, cfg, map[string]Backend{"default": backend}, "default")
}

func NewManagerWithBackends(logger *logrus.Logger, cfg *config.StorageProxyConfig, backends map[string]Backend, defaultBackend string) *Manager {
	if backends == nil {
		backends = make(map[string]Backend)
	}
	return &Manager{
		volumes:        make(map[string]*VolumeContext),
		mountSessions:  make(map[string]map[string]*MountSession),
		directMounts:   make(map[string]*directMountLease),
		invalidates:    make(map[string]map[string]*invalidateTracker),
		logger:         logger,
		config:         cfg,
		backends:       backends,
		defaultBackend: defaultBackend,
	}
}

// SetMountRegistrar sets the mount registrar for distributed coordination.
// This should be called after coordinator is initialized.
func (m *Manager) SetMountRegistrar(registrar MountRegistrar) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.registrar = registrar
}

func (m *Manager) SetStorageObserver(observer StorageObserver) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.observer = observer
}

func (m *Manager) SetMetrics(metrics *obsmetrics.StorageProxyMetrics) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.metrics = metrics
}

// MountVolume mounts a volume using the configured storage backend.
// AccessMode is enforced per storage-proxy instance (not per sandbox).
func (m *Manager) MountVolume(ctx context.Context, s3Prefix, volumeID, teamID string, accessMode AccessMode) (string, time.Time, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	accessMode = NormalizeAccessMode(string(accessMode))
	sessionID := uuid.New().String()
	sessionTime := time.Now()

	if teamID == "" {
		return "", time.Time{}, fmt.Errorf("missing team id for volume mount")
	}

	// Check if already mounted
	if existing, exists := m.volumes[volumeID]; exists {
		if existing.TeamID != "" && existing.TeamID != teamID {
			return "", time.Time{}, fmt.Errorf("volume %s already mounted by another team", volumeID)
		}
		if existing.Access != accessMode {
			return "", time.Time{}, fmt.Errorf("volume %s already mounted with access_mode=%s", volumeID, existing.Access)
		}
		if m.mountSessions[volumeID] == nil {
			m.mountSessions[volumeID] = make(map[string]*MountSession)
		}
		m.mountSessions[volumeID][sessionID] = &MountSession{
			ID:        sessionID,
			TeamID:    teamID,
			CreatedAt: sessionTime,
		}
		return sessionID, sessionTime, nil
	}

	m.logger.WithField("volume_id", volumeID).Info("Mounting volume")

	registeredMount := false
	if m.registrar != nil {
		if err := m.registrar.RegisterMount(ctx, volumeID, MountOptions{
			AccessMode: accessMode,
		}); err != nil {
			return "", time.Time{}, err
		}
		registeredMount = true
	}

	backend, err := m.selectBackend()
	if err != nil {
		if registeredMount {
			_ = m.registrar.UnregisterMount(ctx, volumeID)
		}
		return "", time.Time{}, err
	}

	volCtx, err := backend.MountVolume(ctx, BackendMountRequest{
		S3Prefix:        s3Prefix,
		VolumeID:        volumeID,
		TeamID:          teamID,
		AccessMode:      accessMode,
		MountedAt:       sessionTime,
		Metrics:         m.metrics,
		StorageObserver: m.observer,
	})
	if err != nil {
		if registeredMount {
			_ = m.registrar.UnregisterMount(ctx, volumeID)
		}
		return "", time.Time{}, err
	}

	m.volumes[volumeID] = volCtx
	if m.mountSessions[volumeID] == nil {
		m.mountSessions[volumeID] = make(map[string]*MountSession)
	}
	m.mountSessions[volumeID][sessionID] = &MountSession{
		ID:        sessionID,
		TeamID:    teamID,
		CreatedAt: sessionTime,
	}

	m.logger.WithFields(logrus.Fields{
		"volume_id":   volumeID,
		"backend":     volCtx.Backend,
		"cache_dir":   volCtx.CacheDir,
		"access_mode": accessMode,
	}).Info("Volume mounted successfully")

	return sessionID, sessionTime, nil
}

func (m *Manager) selectBackend() (Backend, error) {
	if backend, ok := m.backends[m.defaultBackend]; ok && backend != nil {
		return backend, nil
	}
	if len(m.backends) == 1 {
		for _, backend := range m.backends {
			if backend != nil {
				return backend, nil
			}
		}
	}
	return nil, fmt.Errorf("storage backend %q is not configured", m.defaultBackend)
}

// UnmountVolume unmounts a volume session and unmounts the volume if it is the last session.
func (m *Manager) UnmountVolume(ctx context.Context, volumeID, sessionID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	volCtx, ok := m.volumes[volumeID]
	if !ok {
		return fmt.Errorf("volume %s not mounted", volumeID)
	}
	if sessionID == "" {
		return fmt.Errorf("missing mount session id")
	}
	sessions, ok := m.mountSessions[volumeID]
	if !ok {
		return fmt.Errorf("mount session %s not found for volume %s", sessionID, volumeID)
	}
	if _, ok := sessions[sessionID]; !ok {
		return fmt.Errorf("mount session %s not found for volume %s", sessionID, volumeID)
	}
	delete(sessions, sessionID)
	if lease := m.directMounts[volumeID]; lease != nil && lease.SessionID == sessionID {
		delete(m.directMounts, volumeID)
	}
	m.clearSessionFromInvalidatesLocked(volumeID, sessionID)
	if len(sessions) > 0 {
		return nil
	}
	delete(m.mountSessions, volumeID)

	m.logger.WithField("volume_id", volumeID).Info("Unmounting volume")

	// Unregister mount from distributed coordination (if registrar is set)
	if m.registrar != nil {
		if err := m.registrar.UnregisterMount(ctx, volumeID); err != nil {
			m.logger.WithError(err).Warn("Failed to unregister mount from coordination")
			// Don't fail the unmount operation
		}
	}

	backend, err := m.selectBackend()
	if err != nil {
		m.logger.WithError(err).Warn("Failed to resolve volume backend for unmount")
	} else if err := backend.UnmountVolume(ctx, volCtx); err != nil {
		m.logger.WithError(err).Warn("Backend volume unmount reported errors")
	}

	delete(m.directMounts, volumeID)
	delete(m.volumes, volumeID)

	m.logger.WithField("volume_id", volumeID).Info("Volume unmounted successfully")

	return nil
}

// GetVolume retrieves volume context
func (m *Manager) GetVolume(volumeID string) (*VolumeContext, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	volCtx, ok := m.volumes[volumeID]
	if !ok {
		return nil, fmt.Errorf("volume %s not mounted", volumeID)
	}

	return volCtx, nil
}

// BeginInvalidate registers an invalidate event that requires remount acks.
func (m *Manager) BeginInvalidate(volumeID, invalidateID string) (int, error) {
	if invalidateID == "" {
		return 0, fmt.Errorf("missing invalidate id")
	}
	m.mu.Lock()
	defer m.mu.Unlock()

	sessions := m.mountSessions[volumeID]
	if len(sessions) == 0 {
		return 0, nil
	}
	if m.invalidates[volumeID] == nil {
		m.invalidates[volumeID] = make(map[string]*invalidateTracker)
	}
	if _, exists := m.invalidates[volumeID][invalidateID]; exists {
		return 0, fmt.Errorf("invalidate %s already registered", invalidateID)
	}
	pending := make(map[string]struct{}, len(sessions))
	for sessionID, session := range sessions {
		if session != nil && session.Scope == MountSessionScopeDirect {
			continue
		}
		pending[sessionID] = struct{}{}
	}
	m.invalidates[volumeID][invalidateID] = &invalidateTracker{
		pending: pending,
		done:    make(chan struct{}),
	}
	return len(pending), nil
}

// WaitForInvalidate waits until all sessions acknowledge the invalidate event.
func (m *Manager) WaitForInvalidate(ctx context.Context, volumeID, invalidateID string) error {
	m.mu.RLock()
	tracker := m.invalidates[volumeID][invalidateID]
	m.mu.RUnlock()
	if tracker == nil {
		return nil
	}

	select {
	case <-tracker.done:
	case <-ctx.Done():
		return ctx.Err()
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	tracker = m.invalidates[volumeID][invalidateID]
	if tracker == nil {
		return nil
	}
	err := tracker.err
	delete(m.invalidates[volumeID], invalidateID)
	if len(m.invalidates[volumeID]) == 0 {
		delete(m.invalidates, volumeID)
	}
	return err
}

// AckInvalidate acknowledges an invalidate event for a mount session.
func (m *Manager) AckInvalidate(volumeID, sessionID, invalidateID string, success bool, errorMessage string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	tracker := m.invalidates[volumeID][invalidateID]
	if tracker == nil {
		return nil
	}
	if _, ok := tracker.pending[sessionID]; !ok {
		return nil
	}
	delete(tracker.pending, sessionID)
	if !success && tracker.err == nil {
		if errorMessage == "" {
			errorMessage = "remount failed"
		}
		tracker.err = fmt.Errorf("%s", errorMessage)
	}
	if tracker.err != nil || len(tracker.pending) == 0 {
		select {
		case <-tracker.done:
		default:
			close(tracker.done)
		}
	}
	return nil
}

func (m *Manager) clearSessionFromInvalidatesLocked(volumeID, sessionID string) {
	trackerMap := m.invalidates[volumeID]
	if len(trackerMap) == 0 {
		return
	}
	for _, tracker := range trackerMap {
		if _, ok := tracker.pending[sessionID]; !ok {
			continue
		}
		delete(tracker.pending, sessionID)
		if tracker.err != nil || len(tracker.pending) == 0 {
			select {
			case <-tracker.done:
			default:
				close(tracker.done)
			}
		}
	}
}

// UpdateVolumeRoot updates the root inode for a mounted volume.
func (m *Manager) UpdateVolumeRoot(volumeID string, rootInode fsmeta.Ino) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	volCtx, ok := m.volumes[volumeID]
	if !ok {
		return fmt.Errorf("volume %s not mounted", volumeID)
	}
	volCtx.RootInode = rootInode
	return nil
}

// ListVolumes returns all mounted volumes
func (m *Manager) ListVolumes() []string {
	m.mu.RLock()
	defer m.mu.RUnlock()

	volumes := make([]string, 0, len(m.volumes))
	for volumeID := range m.volumes {
		volumes = append(volumes, volumeID)
	}

	return volumes
}

// ListMountSessions returns all mount session IDs for a volume.
func (m *Manager) ListMountSessions(volumeID string) []string {
	m.mu.RLock()
	defer m.mu.RUnlock()

	sessions := m.mountSessions[volumeID]
	if len(sessions) == 0 {
		return nil
	}
	ids := make([]string, 0, len(sessions))
	for sessionID := range sessions {
		ids = append(ids, sessionID)
	}
	return ids
}

// AcquireDirectVolumeFileMount acquires a shared direct session for HTTP volume file APIs.
func (m *Manager) AcquireDirectVolumeFileMount(ctx context.Context, volumeID string, mountFn func(context.Context) (string, error)) (func(), error) {
	if volumeID == "" {
		return func() {}, fmt.Errorf("missing volume id")
	}
	if mountFn == nil {
		return func() {}, fmt.Errorf("missing direct mount function")
	}

	now := time.Now()

	m.mu.Lock()
	if lease := m.directMounts[volumeID]; lease != nil {
		if sessions := m.mountSessions[volumeID]; sessions != nil {
			if session := sessions[lease.SessionID]; session != nil {
				session.Scope = MountSessionScopeDirect
				lease.InFlight++
				lease.LastUsed = now
				m.mu.Unlock()
				return m.releaseDirectVolumeFileMount(volumeID, lease.SessionID), nil
			}
		}
		delete(m.directMounts, volumeID)
	}
	m.mu.Unlock()

	sessionID, err := mountFn(ctx)
	if err != nil {
		return func() {}, err
	}

	redundantSessionID := ""

	m.mu.Lock()
	if lease := m.directMounts[volumeID]; lease != nil {
		if sessions := m.mountSessions[volumeID]; sessions != nil {
			if session := sessions[lease.SessionID]; session != nil {
				session.Scope = MountSessionScopeDirect
				lease.InFlight++
				lease.LastUsed = now
				redundantSessionID = sessionID
				m.mu.Unlock()
				if err := m.UnmountVolume(context.Background(), volumeID, redundantSessionID); err != nil {
					m.logger.WithError(err).WithFields(logrus.Fields{
						"volume_id":  volumeID,
						"session_id": redundantSessionID,
					}).Warn("Failed to cleanup redundant direct volume session")
				}
				return m.releaseDirectVolumeFileMount(volumeID, lease.SessionID), nil
			}
		}
		delete(m.directMounts, volumeID)
	}

	if sessions := m.mountSessions[volumeID]; sessions != nil {
		if session := sessions[sessionID]; session != nil {
			session.Scope = MountSessionScopeDirect
		}
	}
	m.directMounts[volumeID] = &directMountLease{
		SessionID: sessionID,
		InFlight:  1,
		LastUsed:  now,
	}
	m.mu.Unlock()

	return m.releaseDirectVolumeFileMount(volumeID, sessionID), nil
}

func (m *Manager) RefreshDirectVolumeFileMount(ctx context.Context, volumeID string) error {
	volCtx, err := m.GetVolume(volumeID)
	if err != nil || volCtx == nil || volCtx.S0FS == nil {
		return err
	}
	_, err = volCtx.S0FS.RefreshMaterialized(ctx)
	return err
}

func (m *Manager) SyncDirectVolumeFileMount(ctx context.Context, volumeID string) error {
	m.mu.RLock()
	_, hasDirectLease := m.directMounts[volumeID]
	_, hasLocalMount := m.volumes[volumeID]
	m.mu.RUnlock()
	if !hasDirectLease && !hasLocalMount {
		return nil
	}
	volCtx, err := m.GetVolume(volumeID)
	if err != nil || volCtx == nil || volCtx.S0FS == nil {
		return err
	}
	manifest, err := volCtx.S0FS.SyncMaterialize(ctx)
	if err != nil {
		return err
	}
	return m.observeMaterializedManifest(ctx, volCtx, manifest)
}

func (m *Manager) observeMaterializedManifest(ctx context.Context, volCtx *VolumeContext, manifest *s0fs.Manifest) error {
	if volCtx == nil || volCtx.Observer == nil || manifest == nil || manifest.State == nil {
		return nil
	}
	if err := volCtx.Observer.ObserveVolumeState(ctx, volCtx.VolumeID, volCtx.TeamID, manifest.State, time.Now().UTC()); err != nil {
		m.logger.WithError(err).WithField("volume_id", volCtx.VolumeID).Warn("Failed to record volume storage observation")
		return err
	}
	return nil
}

func (m *Manager) releaseDirectVolumeFileMount(volumeID, sessionID string) func() {
	return func() {
		m.mu.Lock()
		defer m.mu.Unlock()

		lease := m.directMounts[volumeID]
		if lease == nil || lease.SessionID != sessionID {
			return
		}
		if lease.InFlight > 0 {
			lease.InFlight--
		}
		lease.LastUsed = time.Now()
	}
}

// CleanupIdleDirectVolumeFileMount releases the shared direct session for a volume when no requests are in flight.
func (m *Manager) CleanupIdleDirectVolumeFileMount(ctx context.Context, volumeID string) (bool, error) {
	if volumeID == "" {
		return false, fmt.Errorf("missing volume id")
	}

	m.mu.Lock()
	lease := m.directMounts[volumeID]
	if lease == nil {
		m.mu.Unlock()
		return false, nil
	}
	if lease.InFlight > 0 {
		m.mu.Unlock()
		return false, nil
	}
	sessionID := lease.SessionID
	delete(m.directMounts, volumeID)
	m.mu.Unlock()

	if err := m.UnmountVolume(ctx, volumeID, sessionID); err != nil {
		return false, err
	}
	return true, nil
}

// CleanupIdleDirectVolumeFileMounts unmounts idle shared direct sessions after the idle TTL elapses.
func (m *Manager) CleanupIdleDirectVolumeFileMounts(ctx context.Context, idleTTL time.Duration) []error {
	if idleTTL <= 0 {
		return nil
	}

	type pendingUnmount struct {
		volumeID  string
		sessionID string
	}
	candidates := make([]pendingUnmount, 0)
	cutoff := time.Now().Add(-idleTTL)

	m.mu.Lock()
	for volumeID, lease := range m.directMounts {
		if lease == nil || lease.InFlight > 0 || lease.LastUsed.After(cutoff) {
			continue
		}
		candidates = append(candidates, pendingUnmount{
			volumeID:  volumeID,
			sessionID: lease.SessionID,
		})
		delete(m.directMounts, volumeID)
	}
	m.mu.Unlock()

	var errs []error
	for _, candidate := range candidates {
		if err := m.UnmountVolume(ctx, candidate.volumeID, candidate.sessionID); err != nil {
			errs = append(errs, err)
			m.logger.WithError(err).WithFields(logrus.Fields{
				"volume_id":  candidate.volumeID,
				"session_id": candidate.sessionID,
			}).Warn("Failed to cleanup idle direct volume file mount")
		}
	}
	return errs
}
