package mgocrud

import (
	"errors"

	mgo "gopkg.in/mgo.v2"
)

var (
	ErrNotFound = errors.New("not found")
)

type Session struct {
	session *mgo.Session
}

func (s *Session) Close() {
	s.session.Close()
}

func (s *Session) Copy() *Session {
	return &Session{session: s.session.Copy()}
}

func NewSession(dial string) (*Session, error) {
	session, err := mgo.Dial(dial)
	if err != nil {
		return nil, err
	}
	session.SetMode(mgo.Monotonic, true)
	return &Session{session: session}, nil
}

type Database struct {
	database *mgo.Database
	session  *Session
}

func (s *Session) DB(name string) *Database {
	return &Database{database: s.session.DB(name), session: s}
}

type Collection struct {
	collection *mgo.Collection
}

func (db *Database) Session() *Session {
	return db.session
}

func (db *Database) C(name string) *Collection {
	return &Collection{collection: db.database.C(name)}
}

func (db *Database) Name() string {
	return db.database.Name
}

func (c *Collection) Insert(docs ...interface{}) error {
	return c.collection.Insert(docs...)
}

func (c *Collection) UpdateId(id interface{}, update interface{}) error {
	return c.collection.UpdateId(id, update)
}

func (c *Collection) RemoveId(id interface{}) error {
	return c.collection.RemoveId(id)
}

type ChangeInfo struct {
	changeInfo *mgo.ChangeInfo
}

func (ci *ChangeInfo) Matched() int {
	return ci.changeInfo.Matched
}

func (ci *ChangeInfo) Removed() int {
	return ci.changeInfo.Removed
}

func (ci *ChangeInfo) Updated() int {
	return ci.changeInfo.Updated
}

func (c *Collection) Upsert(selector interface{}, update interface{}) (*ChangeInfo, error) {
	ci, err := c.collection.Upsert(selector, update)
	if err != nil {
		return nil, err
	}
	return &ChangeInfo{changeInfo: ci}, nil
}

func (c *Collection) RemoveAll(filter interface{}) (*ChangeInfo, error) {
	ci, err := c.collection.RemoveAll(filter)
	if err != nil {
		return nil, err
	}
	return &ChangeInfo{changeInfo: ci}, nil
}

type Query struct {
	query *mgo.Query
}

func (c *Collection) FindId(id interface{}) *Query {
	return &Query{query: c.collection.FindId(id)}
}

func (c *Collection) Find(query interface{}) *Query {
	return &Query{query: c.collection.Find(query)}
}

func (q *Query) Select(selector interface{}) *Query {
	q.query = q.query.Select(selector)
	return q
}

func (q *Query) One(result interface{}) error {
	err := q.query.One(result)
	if err == mgo.ErrNotFound {
		err = ErrNotFound
	}
	return err
}

func (q *Query) Sort(fields ...string) *Query {
	q.query = q.query.Sort(fields...)
	return q
}

func (q *Query) Skip(n int) *Query {
	q.query = q.query.Skip(n)
	return q
}

func (q *Query) Limit(n int) *Query {
	q.query = q.query.Limit(n)
	return q
}

func (q *Query) All(result interface{}) error {
	err := q.query.All(result)
	if err == mgo.ErrNotFound {
		err = ErrNotFound
	}
	return err
}

func (q *Query) Count() (int, error) {
	c, err := q.query.Count()
	if err == mgo.ErrNotFound {
		err = ErrNotFound
	}
	return c, err
}

type Index struct {
	index *mgo.Index
}

func NewMgoIndex(index mgo.Index) *Index {
	return &Index{index: &index}
}

func (c *Collection) EnsureIndex(index *Index) error {
	if index != nil && index.index != nil {
		return c.collection.EnsureIndex(*index.index)
	}
	return errors.New("index parameter not initialized with mgo.Index")
}

type Pipe struct {
	pipe *mgo.Pipe
}

func (p *Pipe) All(result interface{}) error {
	err := p.pipe.All(result)
	if err == mgo.ErrNotFound {
		err = ErrNotFound
	}
	return err
}

func (c *Collection) Pipe(pipeline interface{}) *Pipe {
	return &Pipe{pipe: c.collection.Pipe(pipeline).AllowDiskUse()}
}