package handlers

import "errors"
import "io"
import "sync"

import "golang.org/x/crypto/ssh"
import "golang.org/x/net/websocket"

type sshActiveSession struct {
	mu          sync.Mutex
	ws          *websocket.Conn
	client      *ssh.Client
	session     *ssh.Session
	stdin       io.WriteCloser
	closeReason string
	closed      bool
}

type SSHSessionRegistry struct {
	mu    sync.Mutex
	items map[string]*sshActiveSession
}

func NewSSHSessionRegistry() *SSHSessionRegistry {
	var registry *SSHSessionRegistry

	registry = &SSHSessionRegistry{
		items: map[string]*sshActiveSession{},
	}
	return registry
}

func (r *SSHSessionRegistry) Register(id string, item *sshActiveSession) error {
	if r == nil || item == nil {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.items[id] != nil {
		return errors.New("ssh session already active")
	}
	r.items[id] = item
	return nil
}

func (r *SSHSessionRegistry) Unregister(id string, item *sshActiveSession) {
	if r == nil || item == nil {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.items[id] == item {
		delete(r.items, id)
	}
}

func (r *SSHSessionRegistry) RequestClose(id string, reason string) bool {
	var item *sshActiveSession

	if r == nil {
		return false
	}
	r.mu.Lock()
	item = r.items[id]
	r.mu.Unlock()
	if item == nil {
		return false
	}
	item.RequestClose(reason)
	return true
}

func (s *sshActiveSession) SetResources(ws *websocket.Conn, client *ssh.Client, session *ssh.Session, stdin io.WriteCloser) {
	if s == nil {
		return
	}
	s.mu.Lock()
	if ws != nil {
		s.ws = ws
	}
	if client != nil {
		s.client = client
	}
	if session != nil {
		s.session = session
	}
	if stdin != nil {
		s.stdin = stdin
	}
	s.mu.Unlock()
}

func (s *sshActiveSession) CloseReason() string {
	var reason string

	if s == nil {
		return ""
	}
	s.mu.Lock()
	reason = s.closeReason
	s.mu.Unlock()
	return reason
}

func (s *sshActiveSession) RequestClose(reason string) {
	var ws *websocket.Conn
	var client *ssh.Client
	var session *ssh.Session
	var stdin io.WriteCloser

	if s == nil {
		return
	}
	s.mu.Lock()
	if s.closeReason == "" {
		s.closeReason = reason
	}
	if s.closed {
		s.mu.Unlock()
		return
	}
	s.closed = true
	ws = s.ws
	client = s.client
	session = s.session
	stdin = s.stdin
	s.mu.Unlock()
	if stdin != nil {
		_ = stdin.Close()
	}
	if session != nil {
		_ = session.Close()
	}
	if client != nil {
		_ = client.Close()
	}
	if ws != nil {
		_ = ws.Close()
	}
}