package cql

import (
	"database/sql"
	"fmt"
)

type streamLease struct {
	conn    *sql.Conn
	release func()
}

func (c *Conn) beginStream(stream int16) error {
	if c == nil {
		return fmt.Errorf("connection is nil")
	}
	c.streamMu.Lock()
	defer c.streamMu.Unlock()
	if c.streamLeases == nil {
		c.streamLeases = make(map[int16]*streamLease)
	}
	if _, exists := c.streamLeases[stream]; exists {
		return fmt.Errorf("stream %d is already in use", stream)
	}
	c.streamLeases[stream] = nil
	return nil
}

func (c *Conn) attachStreamLease(stream int16, lease *streamLease) bool {
	if c == nil {
		return false
	}
	c.streamMu.Lock()
	defer c.streamMu.Unlock()
	if c.streamLeases == nil {
		return false
	}
	if _, exists := c.streamLeases[stream]; !exists {
		return false
	}
	c.streamLeases[stream] = lease
	return true
}

func (c *Conn) releaseStream(stream int16) {
	lease := c.popStreamLease(stream)
	if lease == nil {
		return
	}
	if lease.release != nil {
		lease.release()
		return
	}
	if lease.conn != nil {
		_ = lease.conn.Close()
	}
}

func (c *Conn) popStreamLease(stream int16) *streamLease {
	if c == nil {
		return nil
	}
	c.streamMu.Lock()
	lease := c.streamLeases[stream]
	delete(c.streamLeases, stream)
	c.streamMu.Unlock()
	return lease
}

func (c *Conn) releaseAllStreamLeases() {
	if c == nil {
		return
	}
	c.streamMu.Lock()
	leases := make([]*streamLease, 0, len(c.streamLeases))
	for _, lease := range c.streamLeases {
		if lease != nil {
			leases = append(leases, lease)
		}
	}
	c.streamLeases = make(map[int16]*streamLease)
	c.streamMu.Unlock()

	for _, lease := range leases {
		if lease.release != nil {
			lease.release()
			continue
		}
		if lease.conn != nil {
			_ = lease.conn.Close()
		}
	}
}