package session

import (
	"bytes"
	"cmp"
	"crypto/rand"
	"encoding/base32"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/creack/pty"
	xterm "github.com/gitpod-io/xterm-go"
	"github.com/samber/lo"
	"golang.org/x/term"
	"tobi/internal/protocol"
	"tobi/internal/target"
	"tobi/internal/termcolor"
)

var (
	ErrDetached = errors.New("session detached")
	b32         = base32.NewEncoding("abcdefghijklmnopqrstuvwxyz234567").WithPadding(base32.NoPadding)
)

// how long an owner outlives its last viewer; atomic because a test shortens
// it while owners from earlier tests are still running
var detachedTTL atomic.Int64

func init() { detachedTTL.Store(int64(24 * time.Hour)) }

type Session struct {
	ID        string    `json:"id"`
	Host      string    `json:"host"`
	Command   string    `json:"command"`
	Cwd       string    `json:"cwd"`
	StartedAt time.Time `json:"started_at"`
	Attached  bool      `json:"attached,omitempty"`
}

func (s Session) Canonical() string { return s.Host + "." + s.ID }

type HopRequest struct {
	Target string
}

func (h HopRequest) Error() string { return "hop to " + h.Target }

func ValidateID(id string) bool {
	return len(id) >= 1 && len(id) <= 64 && strings.TrimLeft(id, "abcdefghijklmnopqrstuvwxyz0123456789-_") == ""
}

func GenerateID() string {
	b := make([]byte, 10)
	rand.Read(b)
	return b32.EncodeToString(b)
}

func RuntimeDir() string {
	base := filepath.Join(os.TempDir(), fmt.Sprintf("tobi-%d", os.Getuid()))
	if xdg := cmp.Or(os.Getenv("XDG_RUNTIME_DIR"), loginRuntimeDir()); xdg != "" {
		base = filepath.Join(xdg, "tobi")
	}
	if fi, err := os.Lstat(base); err == nil && (fi.Mode()&os.ModeSymlink != 0 || int(fi.Sys().(*syscall.Stat_t).Uid) != os.Getuid()) {
		_ = os.RemoveAll(base)
	}
	_ = os.MkdirAll(base, 0700)
	_ = os.Chmod(base, 0700)
	return base
}

func loginRuntimeDir() string {
	d := fmt.Sprintf("/run/user/%d", os.Getuid())
	if fi, err := os.Stat(d); err == nil && fi.IsDir() && int(fi.Sys().(*syscall.Stat_t).Uid) == os.Getuid() {
		return d
	}
	return ""
}

func SocketPath(id string) string {
	if !ValidateID(id) {
		return ""
	}
	return filepath.Join(RuntimeDir(), id+".sock")
}

func IsLocal(host string) bool {
	return strings.EqualFold(host, ShortHostname()) || strings.EqualFold(host, "localhost")
}

func ShortHostname() string {
	h, _ := os.Hostname()
	h, _, _ = strings.Cut(h, ".")
	return lo.Ternary(h != "", h, "localhost")
}

func EnsureOwner(id string, argv []string) error {
	if !ValidateID(id) {
		return fmt.Errorf("invalid session id: %q", id)
	}
	if len(argv) == 0 {
		return errors.New("empty command")
	}
	sockPath := SocketPath(id)
	if answers(sockPath, 30*time.Millisecond) {
		return nil
	}
	if _, err := exec.LookPath(argv[0]); err != nil {
		return err
	}
	cmd := exec.Command(Executable(), argv...)
	cmd.Env = append(os.Environ(), "TOBI_OWNER=1", "TOBI_ID="+id)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	// the owner is detached, so its stderr is the only way it can explain itself
	var stderr syncBuffer
	cmd.Stderr = &stderr
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("start owner: %w", err)
	}
	exited := make(chan struct{})
	go func() { _ = cmd.Wait(); close(exited) }()

	fail := func(reason string) error {
		if msg := stderr.lastLine(); msg != "" {
			return fmt.Errorf("owner for session %q %s: %s", id, reason, msg)
		}
		return fmt.Errorf("owner for session %q %s", id, reason)
	}
	for deadline := time.Now().Add(2 * time.Second); time.Now().Before(deadline); time.Sleep(10 * time.Millisecond) {
		if answers(sockPath, 20*time.Millisecond) {
			return nil
		}
		select {
		case <-exited:
			return fail("exited before it was ready")
		default:
		}
	}
	return fail("did not start")
}

// Executable is this binary, to re-run as an owner or ship to a remote
func Executable() string {
	if self, err := os.Executable(); err == nil {
		return self
	}
	return os.Args[0]
}

// read while the child may still be writing, so the copier must not be raced
type syncBuffer struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (b *syncBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.buf.Len() > 8<<10 {
		return len(p), nil
	}
	return b.buf.Write(p)
}

func (b *syncBuffer) lastLine() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	lines := strings.Split(strings.TrimSpace(b.buf.String()), "\n")
	return strings.TrimSpace(lines[len(lines)-1])
}

func ListLocal() []Session {
	dir := RuntimeDir()
	entries, _ := os.ReadDir(dir)
	var list []Session

	for _, e := range entries {
		id, ok := strings.CutSuffix(e.Name(), ".sock")
		if !ok || e.IsDir() || !ValidateID(id) {
			continue
		}
		conn, err := net.DialTimeout("unix", filepath.Join(dir, e.Name()), 250*time.Millisecond)
		if err != nil {
			_ = os.Remove(filepath.Join(dir, e.Name()))
			continue
		}
		_ = conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
		_ = protocol.WriteFrame(conn, protocol.TypeInfo, nil)
		t, p, err := protocol.ReadFrame(conn)
		_ = conn.Close()
		var s Session
		if err == nil && t == protocol.TypeInfoResp && json.Unmarshal(p, &s) == nil {
			list = append(list, s)
		}
	}
	return list
}

func Attach(stream io.ReadWriteCloser, inputCh <-chan []byte, colors map[int]termcolor.Color) error {
	defer stream.Close()
	size := func(t protocol.FrameType) error {
		c, r, _ := term.GetSize(int(os.Stdout.Fd()))
		return protocol.WriteFrame(stream, t, append(protocol.EncodeSize(c, r), termcolor.EncodeColors(colors)...))
	}
	if err := size(protocol.TypeAttach); err != nil {
		return err
	}
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGWINCH)
	defer signal.Stop(sigCh)
	done := make(chan struct{})
	defer close(done)
	go func() {
		for {
			select {
			case input, ok := <-inputCh:
				if !ok || protocol.WriteFrame(stream, protocol.TypeInput, input) != nil {
					return
				}
			case <-sigCh:
				if size(protocol.TypeResize) != nil {
					return
				}
			case <-done:
				return
			}
		}
	}()
	for {
		t, payload, err := protocol.ReadFrame(stream)
		if err != nil {
			return ErrDetached
		}

		switch t {
		case protocol.TypeSnapshot, protocol.TypeOutput:
			_, _ = os.Stdout.Write(payload)
		case protocol.TypeHop:
			return HopRequest{Target: string(payload)}
		}
	}
}

func SendHop(controlSock, target string) error {
	conn, err := net.Dial("unix", controlSock)
	if err != nil {
		return err
	}
	defer conn.Close()

	return protocol.WriteFrame(conn, protocol.TypeHop, []byte(target))
}

// an owner that predates TypeKill just hangs up, so one that still answers gets a signal instead
func Kill(id string) error {
	sock := SocketPath(id)
	if sock == "" {
		return fmt.Errorf("invalid session id: %q", id)
	}
	conn, err := net.DialTimeout("unix", sock, 250*time.Millisecond)
	if err != nil {
		return fmt.Errorf("session %q is not running", id)
	}
	pid := peerPID(conn)
	_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
	err = protocol.WriteFrame(conn, protocol.TypeKill, nil)
	_, _ = io.Copy(io.Discard, conn)
	_ = conn.Close()
	switch {
	case err != nil:
		return err
	case !answers(sock, 250*time.Millisecond):
		return nil
	case pid <= 0:
		return fmt.Errorf("session %q outlived the kill and its owner pid is unknown", id)
	}
	if err := syscall.Kill(pid, syscall.SIGTERM); err != nil {
		return fmt.Errorf("signal the owner of %q: %w", id, err)
	}
	for deadline := time.Now().Add(time.Second); time.Now().Before(deadline); time.Sleep(20 * time.Millisecond) {
		if !answers(sock, 250*time.Millisecond) {
			return nil
		}
	}
	return fmt.Errorf("session %q outlived SIGTERM", id)
}

func answers(sock string, timeout time.Duration) bool {
	c, err := net.DialTimeout("unix", sock, timeout)
	if err == nil {
		_ = c.Close()
	}
	return err == nil
}

func peerPID(c net.Conn) int {
	uc, ok := c.(*net.UnixConn)
	if !ok {
		return 0
	}
	raw, err := uc.SyscallConn()
	if err != nil {
		return 0
	}
	pid := 0
	_ = raw.Control(func(fd uintptr) { pid = sockPeerPID(int(fd)) })
	return pid
}

type outMsg struct {
	t protocol.FrameType
	p []byte
}

type viewerClient struct {
	conn   net.Conn
	ch     chan outMsg
	done   chan struct{}
	once   sync.Once
	cols   int
	rows   int
	colors map[int]termcolor.Color
}

// send queues msg, dropping a viewer too slow to keep up
func (vc *viewerClient) send(msg outMsg) {
	select {
	case vc.ch <- msg:
	default:
		vc.Close()
	}
}

func (vc *viewerClient) Close() {
	vc.once.Do(func() {
		close(vc.done)
		_ = vc.conn.Close()
	})
}

func pump(vc *viewerClient, snap []byte) {
	_ = protocol.WriteFrame(vc.conn, protocol.TypeSnapshot, snap)
	for {
		select {
		case <-vc.done:
			return
		case msg := <-vc.ch:
			if protocol.WriteFrame(vc.conn, msg.t, msg.p) != nil {
				vc.Close()
				return
			}
		}
	}
}

type Owner struct {
	sess     Session
	ptyFile  *os.File
	vt       *xterm.Terminal
	addon    *xterm.SerializeAddon
	cmd      *exec.Cmd
	listener net.Listener
	mu       sync.Mutex
	viewers  map[*viewerClient]struct{}
	// the viewer that sent input last, so a paste lands on the terminal it came from
	active atomic.Pointer[viewerClient]
	// the osc color handler reads colorSets while mu is held for vt writes,
	// so it gets its own lock
	colorMu   sync.Mutex
	colorSets map[*viewerClient]map[int]termcolor.Color
	idleTimer *time.Timer
	// set once a hop has moved viewers elsewhere: the shell that asked for it
	// only ran tobi, so the owner goes with its last viewer
	retiring bool
	closed   bool
}

func (o *Owner) applySizeLocked() {
	cols, rows := 0, 0
	for vc := range o.viewers {
		if cols == 0 || vc.cols < cols {
			cols = vc.cols
		}
		if rows == 0 || vc.rows < rows {
			rows = vc.rows
		}
	}
	if cols < 1 || rows < 1 {
		return
	}
	_ = pty.Setsize(o.ptyFile, &pty.Winsize{Cols: uint16(cols), Rows: uint16(rows)})
	o.vt.Resize(cols, rows)
}

func (o *Owner) outputLocked(p []byte) {
	for vc := range o.viewers {
		vc.send(outMsg{t: protocol.TypeOutput, p: p})
	}
}

// activeOutputLocked sends p to the viewer that typed last, or to everyone when
// that viewer is gone or nobody has typed yet
func (o *Owner) activeOutputLocked(p []byte) {
	if vc := o.active.Load(); vc != nil {
		if _, ok := o.viewers[vc]; ok {
			vc.send(outMsg{t: protocol.TypeOutput, p: p})
			return
		}
	}
	o.outputLocked(p)
}

func (o *Owner) hop(p []byte) {
	o.mu.Lock()
	viewers := lo.Keys(o.viewers)
	if t := target.Parse(string(p)); len(viewers) > 0 && !(t.ID == o.sess.ID && IsLocal(t.Host)) {
		o.retiring = true
	}
	o.mu.Unlock()
	for _, vc := range viewers {
		vc.send(outMsg{t: protocol.TypeHop, p: p})
	}
}

func (o *Owner) addViewer(conn net.Conn, cols, rows int, colors map[int]termcolor.Color) (*viewerClient, []byte) {
	vc := &viewerClient{conn: conn, ch: make(chan outMsg, 128), done: make(chan struct{}), cols: cols, rows: rows, colors: colors}
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.closed {
		return nil, nil
	}
	if o.idleTimer != nil {
		o.idleTimer.Stop()
		o.idleTimer = nil
	}
	o.viewers[vc] = struct{}{}
	o.colorMu.Lock()
	o.colorSets[vc] = colors
	o.colorMu.Unlock()
	o.applySizeLocked()
	return vc, append(o.addon.Serialize(nil), mouseEncodings[o.vt.DecPrivateModes().MouseEncoding]...)
}

// the serializer restores mouse tracking but not its encoding, so a new viewer would report in X10
var mouseEncodings = map[string]string{"SGR": "\x1b[?1006h", "SGR_PIXELS": "\x1b[?1016h"}

func (o *Owner) removeViewer(vc *viewerClient) {
	o.mu.Lock()
	defer o.mu.Unlock()
	delete(o.viewers, vc)
	o.colorMu.Lock()
	delete(o.colorSets, vc)
	o.colorMu.Unlock()
	switch {
	case len(o.viewers) > 0 || o.closed:
	case o.retiring:
		go o.Close()
	default:
		o.idleTimer = time.AfterFunc(time.Duration(detachedTTL.Load()), o.Close)
	}
	o.applySizeLocked()
}

var csiQueries = map[string]struct{}{
	"\x1b[?6n": {}, "\x1b[6n": {}, "\x1b[=0c": {}, "\x1b[=c": {},
	"\x1b[>0c": {}, "\x1b[>c": {}, "\x1b[0c": {}, "\x1b[c": {},
	"\x1b[?2026$p": {}, "\x1b[>0q": {}, "\x1b[>q": {}, "\x1b[?u": {},
}

func terminalQuery(seq []byte) bool {
	if len(seq) < 3 {
		return false
	}
	if seq[1] == '[' {
		_, ok := csiQueries[string(seq)]
		return ok
	}
	if seq[1] != ']' {
		return false
	}
	end := len(seq) - 1
	if seq[end] == '\a' {
		seq = seq[2:end]
	} else {
		seq = seq[2 : end-1] // querySpan only calls us for a complete ST
	}
	s := string(seq)
	if s == "10;?" || s == "11;?" || s == "12;?" {
		return true
	}
	if !strings.HasPrefix(s, "4;") {
		return false
	}
	parts := strings.Split(s[2:], ";")
	if len(parts) < 2 || len(parts)%2 != 0 {
		return false
	}
	for i := 1; i < len(parts); i += 2 {
		if parts[i] != "?" {
			return false
		}
	}
	return true
}

func querySpan(b []byte) (int, bool) {
	if len(b) < 2 {
		return 0, false
	}
	switch b[1] {
	case '[':
		i := 2
		for i < len(b) && b[i] >= 0x20 && b[i] <= 0x3f {
			i++
		}
		if i == len(b) {
			return 0, false
		}
		if b[i] < 0x40 || b[i] > 0x7e {
			return 1, false
		}
		return i + 1, terminalQuery(b[:i+1])
	case ']':
		for i := 2; i < len(b); i++ {
			if b[i] == '\a' {
				return i + 1, terminalQuery(b[:i+1])
			}
			if b[i] != '\x1b' {
				continue
			}
			if i+1 == len(b) {
				return 0, false
			}
			if b[i+1] == '\\' {
				return i + 2, terminalQuery(b[:i+2])
			}
			return i, true
		}
	default:
		return 2, false
	}
	return 0, false
}

type queryFilter struct{ held []byte }

// what one chunk of pty output becomes
type filtered struct {
	emulator []byte // the owner's own terminal state, which answers the queries
	shared   []byte // every viewer
	routed   []byte // only the active viewer
}

const clipboardOSC = "\x1b]5522;"

// clipboardProbe asks whether the terminal speaks the kitty clipboard protocol.
const clipboardProbe = "\x1b[?5522$p"

// routedToActive reports whether seq belongs to the terminal the user is typing
// into. a clipboard (osc 5522) request prompts there and the reply goes into
// the pty, which two replying terminals would corrupt. the probe is that
// terminal's to answer, not the emulator's, which would report the mode unknown.
func routedToActive(seq []byte) bool {
	return bytes.HasPrefix(seq, []byte(clipboardOSC)) || string(seq) == clipboardProbe
}

// Feed splits pty output. terminal queries stay out of shared because the
// emulator answers them; routed sequences stay out of the emulator.
func (f *queryFilter) Feed(chunk []byte) (out filtered) {
	f.held = append(f.held, chunk...)
	for {
		i := bytes.IndexByte(f.held, 0x1b)
		if i < 0 {
			out.emulator = append(out.emulator, f.held...)
			out.shared = append(out.shared, f.held...)
			f.held = f.held[:0]
			return out
		}
		out.emulator = append(out.emulator, f.held[:i]...)
		out.shared = append(out.shared, f.held[:i]...)
		f.held = f.held[i:]
		n, query := querySpan(f.held)
		if n == 0 && len(f.held) <= 64<<10 {
			return out
		}
		if n == 0 {
			n, query = len(f.held), true
		}
		seq := f.held[:n]
		switch {
		case query:
			out.emulator = append(out.emulator, seq...)
		case routedToActive(seq):
			out.routed = append(out.routed, seq...)
		default:
			out.emulator = append(out.emulator, seq...)
			out.shared = append(out.shared, seq...)
		}
		f.held = f.held[n:]
	}
}

func (o *Owner) readPTYLoop() {
	var filter queryFilter
	buf := make([]byte, 4096)
	for {
		n, err := o.ptyFile.Read(buf)
		if n > 0 {
			o.mu.Lock()
			out := filter.Feed(buf[:n])
			_, _ = o.vt.Write(out.emulator)
			if len(out.shared) > 0 {
				o.outputLocked(out.shared)
			}
			if len(out.routed) > 0 {
				o.activeOutputLocked(out.routed)
			}
			o.mu.Unlock()
		}

		if err != nil {
			o.Close()
			return
		}
	}
}

func (o *Owner) handleConn(conn net.Conn) {
	t, payload, err := protocol.ReadFrame(conn)
	if err != nil {
		_ = conn.Close()
		return
	}
	if t != protocol.TypeAttach {
		switch t {
		case protocol.TypeInfo:
			o.mu.Lock()
			s := o.sess
			s.Attached = len(o.viewers) > 0
			data, _ := json.Marshal(s)
			o.mu.Unlock()
			_ = protocol.WriteFrame(conn, protocol.TypeInfoResp, data)
		case protocol.TypeHop:
			o.hop(payload)
		case protocol.TypeKill:
			o.Close()
		}
		_ = conn.Close()
		return
	}

	cols, rows, _ := protocol.DecodeSize(payload)
	colors, _ := termcolor.DecodeColors(payload[min(4, len(payload)):])
	vc, snap := o.addViewer(conn, cmp.Or(cols, 80), cmp.Or(rows, 24), colors)
	if vc == nil {
		_ = conn.Close()
		return
	}
	go pump(vc, snap)
	for {
		ft, fp, err := protocol.ReadFrame(conn)
		if err != nil {
			break
		}

		switch ft {
		case protocol.TypeInput:
			o.active.Store(vc)
			_, _ = o.ptyFile.Write(fp)
		case protocol.TypeResize:
			if c, r, err := protocol.DecodeSize(fp); err == nil && c >= 1 && r >= 1 {
				o.mu.Lock()
				vc.cols, vc.rows = c, r
				o.applySizeLocked()
				o.mu.Unlock()
			}
		}
	}
	o.removeViewer(vc)
	vc.Close()
}

func (o *Owner) Close() {
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.closed {
		return
	}
	o.closed = true

	if o.idleTimer != nil {
		o.idleTimer.Stop()
	}
	for vc := range o.viewers {
		vc.Close()
	}
	_ = o.listener.Close()
	_ = o.ptyFile.Close()
	_ = o.cmd.Process.Kill()
	o.vt.Dispose()
	_ = os.Remove(o.listener.Addr().String())
}

func RunOwner(id string, command []string) error {
	if !ValidateID(id) || len(command) == 0 {
		return errors.New("invalid id or command")
	}

	sockPath := SocketPath(id)
	_ = os.Remove(sockPath)
	listener, err := net.Listen("unix", sockPath)
	if err != nil {
		return err
	}
	_ = os.Chmod(sockPath, 0700)

	vt := xterm.New(xterm.WithCols(80), xterm.WithRows(24), xterm.WithScrollback(5000))
	cmd := exec.Command(command[0], command[1:]...)
	env := lo.Reject(os.Environ(), func(kv string, _ int) bool {
		return strings.HasPrefix(kv, "CI=")
	})
	if t := os.Getenv("TERM"); t == "" || t == "dumb" {
		env = append(env, "TERM=xterm-256color", "COLORTERM=truecolor")
	}
	cmd.Env = append(env, "TOBI_SESSION="+id, "TOBI_CONTROL="+sockPath)

	ptyFile, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: 80, Rows: 24})
	if err != nil {
		_ = listener.Close()
		_ = os.Remove(sockPath)
		return err
	}

	cwd, _ := os.Getwd()
	owner := &Owner{
		sess: Session{
			ID:        id,
			Host:      ShortHostname(),
			Command:   strings.Join(command, " "),
			Cwd:       cwd,
			StartedAt: time.Now(),
		},
		ptyFile:   ptyFile,
		vt:        vt,
		addon:     xterm.NewSerializeAddon(vt),
		cmd:       cmd,
		listener:  listener,
		viewers:   map[*viewerClient]struct{}{},
		colorSets: map[*viewerClient]map[int]termcolor.Color{},
	}

	vt.OnTitleChange(func(title string) {
		if title != "" {
			owner.sess.Command = title
		}
	})
	vt.RegisterOscHandler(7, xterm.NewOscStringHandler(func(payload string) bool {
		if u, err := url.Parse(payload); err == nil && u.Scheme == "file" && u.Path != "" {
			owner.sess.Cwd = u.Path
		}
		return true
	}))
	vt.OnData(func(s string) {
		_, _ = owner.ptyFile.Write([]byte(s))
	})
	vt.OnColor(func(events []xterm.ColorEvent) {
		for _, ev := range events {
			if ev.Type != xterm.ColorRequestReport || ev.Index < 0 || ev.Index > int(xterm.SpecialColorCursor) {
				continue
			}
			owner.colorMu.Lock()
			sets := lo.Values(owner.colorSets)
			owner.colorMu.Unlock()
			if c, ok := termcolor.Agreed(ev.Index, sets); ok {
				fmt.Fprint(owner.ptyFile, termcolor.Reply(ev.Index, c))
			}
		}
	})
	owner.idleTimer = time.AfterFunc(time.Duration(detachedTTL.Load()), owner.Close)

	go owner.readPTYLoop()
	go func() { _ = cmd.Wait() }()

	for {
		conn, err := listener.Accept()
		if err != nil {
			return nil
		}
		go owner.handleConn(conn)
	}
}