package mongodbutil

import (
	"log"
	"sync"
	"time"

	mgo "gopkg.in/mgo.v2"
)

var Mgo *Pool

type Pool struct {
	mu      sync.RWMutex
	initCap int
	maxCap  int
	timeout int64
	ch      chan *mgosess
	addr    string
	DB      string
	live    int
}


type mgosess struct {
	sess      *mgo.Session
	timestamp int64
}

func MgoFactory(initCap int, maxCap int, timeout int64, addr, DB string) *Pool {
	p := &Pool{sync.RWMutex{}, initCap, maxCap, timeout, make(chan *mgosess, maxCap), addr, DB, 0}
	p.init()
	go p.gc()
	return p
}

func (p *Pool) GetLive() int {
	return p.live
}

func (p *Pool) init() {
	for i := 0; i < p.initCap; i++ {
		sess, err := mgo.DialWithTimeout(p.addr, time.Duration(p.timeout)*time.Second)
		if sess != nil && sess.Ping() == nil {
			p.live++
			p.ch <- &mgosess{sess, time.Now().Unix()}
		} else {
			log.Println(err.Error())
		}
	}
}

func (p *Pool) Get() (sess *mgo.Session) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.live > 0 {
		select {
		case mgos := <-p.ch:
			if mgos.sess.Ping() == nil {
				sess = mgos.sess
			} else {
				p.live--
			}
		case <-time.After(100 * time.Millisecond):
		}
	}
	if sess == nil && p.live < p.maxCap {
		s, err := mgo.DialWithTimeout(p.addr, 10*time.Second)
		if s != nil && s.Ping() == nil {
			p.live++
			sess = s
		} else {
			log.Println(err.Error())
		}
	}
	return
}

func (p *Pool) Close(sess *mgo.Session) {
	if sess != nil {
		if sess.Ping() == nil {
			p.mu.Lock()
			select {
			case p.ch <- &mgosess{sess, time.Now().Unix()}:
			default:
				p.live--
			}
			p.mu.Unlock()
		} else {
			p.live--
		}
	}
}

func (p *Pool) gc() {
	p.mu.Lock()
	size := len(p.ch)
	if size > p.initCap {
		tn := time.Now().Unix()
		init1 := 0
		for i := 0; i < size; i++ {
			select {
			case c := <-p.ch:
				if tn-c.timestamp < p.timeout || init1 < p.initCap {
					p.ch <- c
					init1++
				} else {
					c.sess.Close()
					p.live--
				}
			default:
			}
		}
	}
	p.mu.Unlock()
	//log.Println("size:", size, "live:", p.live)
	time.AfterFunc(time.Duration(p.timeout)*time.Second, p.gc)
}
