package ssh

import (
	"bytes"
	"fmt"
	"net/url"
	"sync"
)

// CWDTracker parses OSC 7 escape sequences from terminal output to track the
// remote working directory.
type CWDTracker struct {
	currentPath string
	mu          sync.RWMutex
}

// NewCWDTracker creates a new CWDTracker.
func NewCWDTracker() *CWDTracker {
	return &CWDTracker{}
}

// osc7Prefix is the escape sequence that starts an OSC 7 directive.
var osc7Prefix = []byte("\033]7;")

// stTerminator is the ST (String Terminator) escape: ESC + backslash.
var stTerminator = []byte("\033\\")

// belTerminator is the BEL character, an alternative OSC terminator.
var belTerminator = []byte{0x07}

// ProcessOutput scans data for OSC 7 sequences of the form:
//
//	\033]7;file://hostname/path\033\\   (ST terminator)
//	\033]7;file://hostname/path\007     (BEL terminator)
//
// It returns cleaned output with all OSC 7 sequences stripped and the new CWD
// path (or "" if no OSC 7 was found).
func (t *CWDTracker) ProcessOutput(data []byte) (cleaned []byte, newCWD string) {
	var result []byte
	remaining := data
	var lastCWD string

	for {
		idx := bytes.Index(remaining, osc7Prefix)
		if idx == -1 {
			result = append(result, remaining...)
			break
		}

		// Append everything before the OSC 7 sequence.
		result = append(result, remaining[:idx]...)

		// Find the end of the OSC 7 payload (after the prefix).
		afterPrefix := remaining[idx+len(osc7Prefix):]

		// Try ST terminator first (\033\\), then BEL (\007).
		endIdx := -1
		terminatorLen := 0

		if stIdx := bytes.Index(afterPrefix, stTerminator); stIdx != -1 {
			endIdx = stIdx
			terminatorLen = len(stTerminator)
		}
		if belIdx := bytes.Index(afterPrefix, belTerminator); belIdx != -1 {
			if endIdx == -1 || belIdx < endIdx {
				endIdx = belIdx
				terminatorLen = len(belTerminator)
			}
		}

		if endIdx == -1 {
			// No terminator found; treat the rest as literal output.
			result = append(result, remaining[idx:]...)
			break
		}

		// Extract the URI payload between prefix and terminator.
		payload := string(afterPrefix[:endIdx])
		if path := extractPathFromOSC7(payload); path != "" {
			lastCWD = path
		}

		remaining = afterPrefix[endIdx+terminatorLen:]
	}

	if lastCWD != "" {
		t.mu.Lock()
		t.currentPath = lastCWD
		t.mu.Unlock()
	}

	return result, lastCWD
}

// GetCWD returns the current tracked working directory.
func (t *CWDTracker) GetCWD() string {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.currentPath
}

// extractPathFromOSC7 parses a file:// URI and returns the path component.
func extractPathFromOSC7(uri string) string {
	u, err := url.Parse(uri)
	if err != nil {
		return ""
	}
	if u.Scheme != "file" {
		return ""
	}
	return u.Path
}

// ShellIntegrationCommand returns the shell command to inject for CWD tracking
// via OSC 7. The returned command sets up a prompt hook that emits the OSC 7
// escape sequence after every command.
func ShellIntegrationCommand(shellType string) string {
	switch shellType {
	case "bash":
		return fmt.Sprintf(`PROMPT_COMMAND='printf "\033]7;file://%%s%%s\033\\" "$HOSTNAME" "$PWD"'`)
	case "zsh":
		return fmt.Sprintf(`precmd() { printf '\033]7;file://%%s%%s\033\\' "$HOST" "$PWD" }`)
	case "fish":
		return `function __wraith_osc7 --on-event fish_prompt; printf '\033]7;file://%s%s\033\\' (hostname) (pwd); end`
	default:
		return ""
	}
}