wraith/internal/app/app.go

package app

import (
	"database/sql"
	"encoding/hex"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"

	"encoding/base64"

	"github.com/vstockwell/wraith/internal/ai"
	"github.com/vstockwell/wraith/internal/connections"
	"github.com/vstockwell/wraith/internal/credentials"
	"github.com/vstockwell/wraith/internal/db"
	"github.com/vstockwell/wraith/internal/importer"
	"github.com/vstockwell/wraith/internal/plugin"
	"github.com/vstockwell/wraith/internal/rdp"
	"github.com/vstockwell/wraith/internal/session"
	"github.com/vstockwell/wraith/internal/settings"
	"github.com/vstockwell/wraith/internal/sftp"
	"github.com/vstockwell/wraith/internal/ssh"
	"github.com/vstockwell/wraith/internal/theme"
	"github.com/vstockwell/wraith/internal/updater"
	"github.com/vstockwell/wraith/internal/vault"
	sftplib "github.com/pkg/sftp"
	gossh "golang.org/x/crypto/ssh"
	"github.com/wailsapp/wails/v3/pkg/application"
)

// WraithApp is the main application struct that wires together all services
// and exposes vault management methods to the frontend via Wails bindings.
type WraithApp struct {
	db          *sql.DB
	Vault       *vault.VaultService
	Settings    *settings.SettingsService
	Connections *connections.ConnectionService
	Themes      *theme.ThemeService
	Sessions    *session.Manager
	Plugins     *plugin.Registry
	SSH         *ssh.SSHService
	SFTP        *sftp.SFTPService
	RDP         *rdp.RDPService
	Credentials *credentials.CredentialService
	AI          *ai.AIService
	Updater     *updater.UpdateService
	oauthMgr    *ai.OAuthManager
	wailsApp    *application.App
	unlocked    bool
}

// SetWailsApp stores the Wails application reference for event emission.
// Must be called after application.New() and before app.Run().
func (a *WraithApp) SetWailsApp(app *application.App) {
	a.wailsApp = app
}

// New creates and initializes the WraithApp, opening the database, running
// migrations, creating all services, and seeding built-in themes.
// The version string is the build-time semver (e.g. "0.2.0") used by the
// auto-updater to check for newer releases.
func New(version string) (*WraithApp, error) {
	dataDir := dataDirectory()
	dbPath := filepath.Join(dataDir, "wraith.db")

	slog.Info("opening database", "path", dbPath)
	database, err := db.Open(dbPath)
	if err != nil {
		return nil, fmt.Errorf("open database: %w", err)
	}

	if err := db.Migrate(database); err != nil {
		return nil, fmt.Errorf("run migrations: %w", err)
	}

	settingsSvc := settings.NewSettingsService(database)
	connSvc := connections.NewConnectionService(database)
	themeSvc := theme.NewThemeService(database)
	sessionMgr := session.NewManager()
	pluginReg := plugin.NewRegistry()

	// SSH output handler — emits Wails events to the frontend.
	// The closure captures `app` (the WraithApp being built). The wailsApp
	// field is set after application.New() in main.go, but SSH sessions only
	// start after app.Run(), so wailsApp is always valid at call time.
	var app *WraithApp
	sshSvc := ssh.NewSSHService(database, func(sessionID string, data []byte) {
		if app != nil && app.wailsApp != nil {
			// Base64 encode binary data for safe transport over Wails events
			app.wailsApp.Event.Emit("ssh:data:"+sessionID, base64.StdEncoding.EncodeToString(data))
		}
	})
	sftpSvc := sftp.NewSFTPService()

	// RDP service with platform-aware backend factory.
	// On Windows the factory returns a FreeRDPBackend backed by libfreerdp3.dll;
	// on other platforms it falls back to MockBackend for development.
	rdpSvc := rdp.NewRDPService(func() rdp.RDPBackend {
		return rdp.NewProductionBackend()
	})

	// CredentialService requires the vault to be unlocked, so it starts nil.
	// It is created lazily after the vault is unlocked via initCredentials().

	// AI copilot services — OAuthManager starts without a vault reference;
	// it will be wired once the vault is unlocked.
	oauthMgr := ai.NewOAuthManager(settingsSvc, nil)
	toolRouter := ai.NewToolRouter()
	toolRouter.SetServices(sshSvc, sftpSvc, rdpSvc, sessionMgr, connSvc)
	convMgr := ai.NewConversationManager(database)
	aiSvc := ai.NewAIService(oauthMgr, toolRouter, convMgr)
	toolRouter.SetAIService(aiSvc)

	// Seed built-in themes on every startup (INSERT OR IGNORE keeps it idempotent)
	if err := themeSvc.SeedBuiltins(); err != nil {
		slog.Warn("failed to seed themes", "error", err)
	}

	updaterSvc := updater.NewUpdateService(version)

	app = &WraithApp{
		db:          database,
		Settings:    settingsSvc,
		Connections: connSvc,
		Themes:      themeSvc,
		Sessions:    sessionMgr,
		Plugins:     pluginReg,
		SSH:         sshSvc,
		SFTP:        sftpSvc,
		RDP:         rdpSvc,
		AI:          aiSvc,
		Updater:     updaterSvc,
		oauthMgr:    oauthMgr,
	}
	return app, nil
}

// dataDirectory returns the path where Wraith stores its data.
// On Windows with APPDATA set, it uses %APPDATA%\Wraith.
// On macOS/Linux with XDG_DATA_HOME or HOME, it uses the appropriate path.
// Falls back to the current working directory for development.
func dataDirectory() string {
	// Windows
	if appData := os.Getenv("APPDATA"); appData != "" {
		return filepath.Join(appData, "Wraith")
	}

	// macOS / Linux: use XDG_DATA_HOME or fallback to ~/.local/share
	if home, err := os.UserHomeDir(); err == nil {
		if xdg := os.Getenv("XDG_DATA_HOME"); xdg != "" {
			return filepath.Join(xdg, "wraith")
		}
		return filepath.Join(home, ".local", "share", "wraith")
	}

	// Dev fallback
	return "."
}

// IsFirstRun checks whether the vault has been set up by looking for vault_salt in settings.
func (a *WraithApp) IsFirstRun() bool {
	salt, _ := a.Settings.Get("vault_salt")
	return salt == ""
}

// CreateVault sets up the vault with a master password. It generates a salt,
// derives an encryption key, and stores the salt and a check value in settings.
func (a *WraithApp) CreateVault(password string) error {
	salt, err := vault.GenerateSalt()
	if err != nil {
		return fmt.Errorf("generate salt: %w", err)
	}

	key := vault.DeriveKey(password, salt)
	a.Vault = vault.NewVaultService(key)

	// Store salt as hex in settings
	if err := a.Settings.Set("vault_salt", hex.EncodeToString(salt)); err != nil {
		return fmt.Errorf("store salt: %w", err)
	}

	// Encrypt a known check value — used to verify the password on unlock
	check, err := a.Vault.Encrypt("wraith-vault-check")
	if err != nil {
		return fmt.Errorf("encrypt check value: %w", err)
	}
	if err := a.Settings.Set("vault_check", check); err != nil {
		return fmt.Errorf("store check value: %w", err)
	}

	a.unlocked = true
	a.initCredentials()
	slog.Info("vault created successfully")
	return nil
}

// Unlock verifies the master password against the stored check value and
// initializes the vault service for decryption.
func (a *WraithApp) Unlock(password string) error {
	saltHex, err := a.Settings.Get("vault_salt")
	if err != nil || saltHex == "" {
		return fmt.Errorf("vault not set up — call CreateVault first")
	}

	salt, err := hex.DecodeString(saltHex)
	if err != nil {
		return fmt.Errorf("decode salt: %w", err)
	}

	key := vault.DeriveKey(password, salt)
	vs := vault.NewVaultService(key)

	// Verify by decrypting the stored check value
	checkEncrypted, err := a.Settings.Get("vault_check")
	if err != nil || checkEncrypted == "" {
		return fmt.Errorf("vault check value missing")
	}

	checkPlain, err := vs.Decrypt(checkEncrypted)
	if err != nil {
		return fmt.Errorf("incorrect master password")
	}
	if checkPlain != "wraith-vault-check" {
		return fmt.Errorf("incorrect master password")
	}

	a.Vault = vs
	a.unlocked = true
	a.initCredentials()
	slog.Info("vault unlocked successfully")
	return nil
}

// IsUnlocked returns whether the vault is currently unlocked.
func (a *WraithApp) IsUnlocked() bool {
	return a.unlocked
}

// initCredentials creates the CredentialService after the vault is unlocked
// and wires the vault to services that need it (OAuth, etc.).
func (a *WraithApp) initCredentials() {
	if a.Vault != nil {
		a.Credentials = credentials.NewCredentialService(a.db, a.Vault)
		if a.oauthMgr != nil {
			a.oauthMgr.SetVault(a.Vault)
		}
	}
}

// ConnectSSH opens an SSH session to the given connection ID.
// It resolves credentials from the vault, builds auth methods, and returns a session ID.
// The frontend calls this instead of SSHService.Connect directly (which takes Go-only types).
func (a *WraithApp) ConnectSSH(connectionID int64, cols, rows int) (string, error) {
	conn, err := a.Connections.GetConnection(connectionID)
	if err != nil {
		return "", fmt.Errorf("connection not found: %w", err)
	}

	// Build SSH auth methods from the connection's credential
	var authMethods []gossh.AuthMethod
	username := "root" // default

	slog.Info("ConnectSSH resolving auth",
		"connectionID", connectionID,
		"host", conn.Hostname,
		"credentialID", conn.CredentialID,
		"vaultUnlocked", a.Credentials != nil,
	)

	if conn.CredentialID != nil && a.Credentials != nil {
		cred, err := a.Credentials.GetCredential(*conn.CredentialID)
		if err != nil {
			slog.Warn("failed to load credential", "id", *conn.CredentialID, "error", err)
		} else {
			slog.Info("credential loaded", "name", cred.Name, "type", cred.Type, "username", cred.Username, "sshKeyID", cred.SSHKeyID)
			if cred.Username != "" {
				username = cred.Username
			}

			switch cred.Type {
			case "password":
				pw, err := a.Credentials.DecryptPassword(cred.ID)
				if err != nil {
					slog.Warn("failed to decrypt password", "error", err)
				} else {
					authMethods = append(authMethods, gossh.Password(pw))
				}
			case "ssh_key":
				if cred.SSHKeyID != nil {
					keyPEM, passphrase, err := a.Credentials.DecryptSSHKey(*cred.SSHKeyID)
					if err != nil {
						slog.Warn("failed to decrypt SSH key", "error", err)
					} else {
						var signer gossh.Signer
						if passphrase != "" {
							signer, err = gossh.ParsePrivateKeyWithPassphrase(keyPEM, []byte(passphrase))
						} else {
							signer, err = gossh.ParsePrivateKey(keyPEM)
						}
						if err != nil {
							slog.Warn("failed to parse SSH key", "error", err)
						} else {
							authMethods = append(authMethods, gossh.PublicKeys(signer))
						}
					}
				}
			}
		}
	}

	// If no stored credentials, return an error telling the frontend to prompt
	if len(authMethods) == 0 {
		return "", fmt.Errorf("NO_CREDENTIALS: no credentials configured for this connection — assign a credential or use ConnectSSHWithPassword")
	}

	sessionID, err := a.SSH.Connect(conn.Hostname, conn.Port, username, authMethods, cols, rows)
	if err != nil {
		return "", fmt.Errorf("SSH connect failed: %w", err)
	}

	// Register SFTP client on the same SSH connection
	if sess, ok := a.SSH.GetSession(sessionID); ok && sess != nil {
		sftpClient, err := sftplib.NewClient(sess.Client)
		if err != nil {
			slog.Warn("failed to create SFTP client", "error", err)
			// Non-fatal — SSH still works without SFTP
		} else {
			a.SFTP.RegisterClient(sessionID, sftpClient)
		}
	}

	// Update last_connected timestamp
	if _, err := a.db.Exec("UPDATE connections SET last_connected = CURRENT_TIMESTAMP WHERE id = ?", connectionID); err != nil {
		slog.Warn("failed to update last_connected", "error", err)
	}

	slog.Info("SSH session started", "sessionID", sessionID, "host", conn.Hostname, "user", username)
	return sessionID, nil
}

// ConnectSSHWithPassword opens an SSH session using an ad-hoc username/password.
// Used when no stored credential exists — the frontend prompts and passes creds directly.
func (a *WraithApp) ConnectSSHWithPassword(connectionID int64, username, password string, cols, rows int) (string, error) {
	conn, err := a.Connections.GetConnection(connectionID)
	if err != nil {
		return "", fmt.Errorf("connection not found: %w", err)
	}

	authMethods := []gossh.AuthMethod{
		gossh.Password(password),
		gossh.KeyboardInteractive(func(name, instruction string, questions []string, echos []bool) ([]string, error) {
			answers := make([]string, len(questions))
			for i := range answers {
				answers[i] = password
			}
			return answers, nil
		}),
	}

	sessionID, err := a.SSH.Connect(conn.Hostname, conn.Port, username, authMethods, cols, rows)
	if err != nil {
		return "", fmt.Errorf("SSH connect failed: %w", err)
	}

	// Register SFTP client
	if sess, ok := a.SSH.GetSession(sessionID); ok && sess != nil {
		sftpClient, err := sftplib.NewClient(sess.Client)
		if err != nil {
			slog.Warn("failed to create SFTP client", "error", err)
		} else {
			a.SFTP.RegisterClient(sessionID, sftpClient)
		}
	}

	if _, err := a.db.Exec("UPDATE connections SET last_connected = CURRENT_TIMESTAMP WHERE id = ?", connectionID); err != nil {
		slog.Warn("failed to update last_connected", "error", err)
	}

	slog.Info("SSH session started (ad-hoc password)", "sessionID", sessionID, "host", conn.Hostname, "user", username)
	return sessionID, nil
}

// GetVersion returns the build-time version string.
func (a *WraithApp) GetVersion() string {
	return a.Updater.CurrentVersion()
}

// DisconnectSession closes an active SSH session and its SFTP client.
func (a *WraithApp) DisconnectSession(sessionID string) error {
	a.SFTP.RemoveClient(sessionID)
	return a.SSH.Disconnect(sessionID)
}

// ConnectRDP opens an RDP session to the given connection ID.
// It resolves credentials from the vault, builds an RDPConfig, and returns a session ID.
// width and height are the initial desktop dimensions in pixels.
func (a *WraithApp) ConnectRDP(connectionID int64, width, height int) (string, error) {
	conn, err := a.Connections.GetConnection(connectionID)
	if err != nil {
		return "", fmt.Errorf("connection not found: %w", err)
	}

	config := rdp.RDPConfig{
		Hostname: conn.Hostname,
		Port:     conn.Port,
		Width:    width,
		Height:   height,
	}

	if conn.CredentialID != nil && a.Credentials != nil {
		cred, err := a.Credentials.GetCredential(*conn.CredentialID)
		if err != nil {
			slog.Warn("failed to load credential", "id", *conn.CredentialID, "error", err)
		} else {
			if cred.Username != "" {
				config.Username = cred.Username
			}
			if cred.Domain != "" {
				config.Domain = cred.Domain
			}

			if cred.Type == "password" {
				pw, err := a.Credentials.DecryptPassword(cred.ID)
				if err != nil {
					slog.Warn("failed to decrypt password", "error", err)
				} else {
					config.Password = pw
				}
			}
		}
	}

	sessionID, err := a.RDP.Connect(config, connectionID)
	if err != nil {
		return "", fmt.Errorf("RDP connect failed: %w", err)
	}

	// Update last_connected timestamp
	if _, err := a.db.Exec("UPDATE connections SET last_connected = CURRENT_TIMESTAMP WHERE id = ?", connectionID); err != nil {
		slog.Warn("failed to update last_connected", "error", err)
	}

	slog.Info("RDP session started", "sessionID", sessionID, "host", conn.Hostname)
	return sessionID, nil
}

// RDPGetFrame returns the current frame for an RDP session as a base64-encoded
// string. Go []byte is serialised by Wails as a base64 string over the JSON
// bridge, so the frontend decodes it with atob() to recover the raw RGBA bytes.
func (a *WraithApp) RDPGetFrame(sessionID string) (string, error) {
	raw, err := a.RDP.GetFrame(sessionID)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(raw), nil
}

// RDPSendMouse forwards a mouse event to an RDP session.
// flags uses the RDP mouse-event flag constants defined in internal/rdp/input.go.
func (a *WraithApp) RDPSendMouse(sessionID string, x, y int, flags uint32) error {
	return a.RDP.SendMouse(sessionID, x, y, flags)
}

// RDPSendKey forwards a key event (scancode + press/release) to an RDP session.
func (a *WraithApp) RDPSendKey(sessionID string, scancode uint32, pressed bool) error {
	return a.RDP.SendKey(sessionID, scancode, pressed)
}

// RDPSendClipboard forwards clipboard text to an RDP session.
func (a *WraithApp) RDPSendClipboard(sessionID string, text string) error {
	return a.RDP.SendClipboard(sessionID, text)
}

// RDPDisconnect tears down an RDP session.
func (a *WraithApp) RDPDisconnect(sessionID string) error {
	return a.RDP.Disconnect(sessionID)
}

// ---------- Credential proxy methods ----------
// CredentialService is nil until the vault is unlocked. These proxies expose
// it via WraithApp (which IS registered as a Wails service at startup).

// ListCredentials returns all stored credentials (no encrypted values).
func (a *WraithApp) ListCredentials() ([]credentials.Credential, error) {
	if a.Credentials == nil {
		return nil, fmt.Errorf("vault is locked")
	}
	return a.Credentials.ListCredentials()
}

// CreatePassword creates a password credential encrypted via the vault.
func (a *WraithApp) CreatePassword(name, username, password, domain string) (*credentials.Credential, error) {
	if a.Credentials == nil {
		return nil, fmt.Errorf("vault is locked")
	}
	return a.Credentials.CreatePassword(name, username, password, domain)
}

// CreateSSHKeyCredential imports an SSH private key and creates a Credential
// record referencing it. privateKeyPEM is the raw PEM string (NOT base64 encoded).
func (a *WraithApp) CreateSSHKeyCredential(name, username string, privateKeyPEM string, passphrase string) (*credentials.Credential, error) {
	if a.Credentials == nil {
		return nil, fmt.Errorf("vault is locked")
	}
	return a.Credentials.CreateSSHKeyCredential(name, username, []byte(privateKeyPEM), passphrase)
}

// DeleteCredential removes a credential by ID.
func (a *WraithApp) DeleteCredential(id int64) error {
	if a.Credentials == nil {
		return fmt.Errorf("vault is locked")
	}
	return a.Credentials.DeleteCredential(id)
}

// ImportMobaConf parses a MobaXTerm .mobaconf file and imports its contents
// (groups, connections, host keys) into the database.
func (a *WraithApp) ImportMobaConf(fileContent string) (*plugin.ImportResult, error) {
	imp := &importer.MobaConfImporter{}
	result, err := imp.Parse([]byte(fileContent))
	if err != nil {
		return nil, fmt.Errorf("parse mobaconf: %w", err)
	}

	// Create groups and track name -> ID mapping
	groupMap := make(map[string]int64)
	for _, g := range result.Groups {
		created, err := a.Connections.CreateGroup(g.Name, nil)
		if err != nil {
			slog.Warn("failed to create group during import", "name", g.Name, "error", err)
			continue
		}
		groupMap[g.Name] = created.ID
	}

	// Create connections
	for _, c := range result.Connections {
		var groupID *int64
		if id, ok := groupMap[c.GroupName]; ok {
			groupID = &id
		}
		_, err := a.Connections.CreateConnection(connections.CreateConnectionInput{
			Name:     c.Name,
			Hostname: c.Hostname,
			Port:     c.Port,
			Protocol: c.Protocol,
			GroupID:  groupID,
			Notes:   c.Notes,
		})
		if err != nil {
			slog.Warn("failed to create connection during import", "name", c.Name, "error", err)
		}
	}

	// Store host keys
	for _, hk := range result.HostKeys {
		_, err := a.db.Exec(
			`INSERT OR IGNORE INTO host_keys (hostname, port, key_type, fingerprint) VALUES (?, ?, ?, ?)`,
			hk.Hostname, hk.Port, hk.KeyType, hk.Fingerprint,
		)
		if err != nil {
			slog.Warn("failed to store host key during import", "hostname", hk.Hostname, "error", err)
		}
	}

	slog.Info("mobaconf import complete",
		"groups", len(result.Groups),
		"connections", len(result.Connections),
		"hostKeys", len(result.HostKeys),
	)
	return result, nil
}