ui_impl.go

package ui

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gen2brain/beeep"
	"github.com/google/uuid"
	"github.com/samber/lo"
	"github.com/tidwall/gjson"
	"time"
	"wox/i18n"
	"wox/plugin"
	"wox/setting"
	"wox/share"
	"wox/util"
	"wox/util/ime"
	"wox/util/screen"
)

type uiImpl struct {
	requestMap *util.HashMap[string, chan WebsocketMsg]
}

func (u *uiImpl) ChangeQuery(ctx context.Context, query share.ChangedQuery) {
	u.send(ctx, "ChangeQuery", query)
}

func (u *uiImpl) HideApp(ctx context.Context) {
	u.send(ctx, "HideApp", nil)
}

func (u *uiImpl) ShowApp(ctx context.Context, showContext share.ShowContext) {
	u.send(ctx, "ShowApp", getShowAppParams(ctx, showContext.SelectAll))
}

func (u *uiImpl) ToggleApp(ctx context.Context) {
	u.send(ctx, "ToggleApp", getShowAppParams(ctx, true))
}

func (u *uiImpl) Notify(ctx context.Context, title string, description string) {
	err := beeep.Notify(title, description, "")
	if err != nil {
		logger.Error(ctx, fmt.Sprintf("notify error: %s", err.Error()))
	}
}

func (u *uiImpl) GetServerPort(ctx context.Context) int {
	return GetUIManager().serverPort
}

func (u *uiImpl) ChangeTheme(ctx context.Context, theme share.Theme) {
	logger.Info(ctx, fmt.Sprintf("change theme: %s", theme.ThemeName))
	woxSetting := setting.GetSettingManager().GetWoxSetting(ctx)
	woxSetting.ThemeId = theme.ThemeId
	setting.GetSettingManager().SaveWoxSetting(ctx)
	u.send(ctx, "ChangeTheme", theme)
}

func (u *uiImpl) OpenSettingWindow(ctx context.Context, windowContext share.SettingWindowContext) {
	u.send(ctx, "OpenSettingWindow", windowContext)
}

func (u *uiImpl) GetAllThemes(ctx context.Context) []share.Theme {
	return GetUIManager().GetAllThemes(ctx)
}

func (u *uiImpl) send(ctx context.Context, method string, data any) (responseData any, responseErr error) {
	requestID := uuid.NewString()
	resultChan := make(chan WebsocketMsg)
	u.requestMap.Store(requestID, resultChan)
	defer u.requestMap.Delete(requestID)

	traceId := util.GetContextTraceId(ctx)

	err := requestUI(ctx, WebsocketMsg{
		RequestId: requestID,
		TraceId:   traceId,
		Method:    method,
		Data:      data,
	})
	if err != nil {
		logger.Error(ctx, fmt.Sprintf("send message to UI error: %s", err.Error()))
		return "", err
	}

	var timeout = time.Second * 2
	if method == "PickFiles" {
		// pick files may take a long time
		timeout = time.Second * 180
	}
	select {
	case <-time.NewTimer(timeout).C:
		logger.Error(ctx, fmt.Sprintf("invoke ui method %s response timeout", method))
		return "", fmt.Errorf("request timeout, request id: %s", requestID)
	case response := <-resultChan:
		if !response.Success {
			return response.Data, errors.New("ui method response error")
		} else {
			return response.Data, nil
		}
	}
}

func (u *uiImpl) PickFiles(ctx context.Context, params share.PickFilesParams) []string {
	respData, err := u.send(ctx, "PickFiles", params)
	if err != nil {
		return nil
	}
	if _, ok := respData.([]any); !ok {
		logger.Error(ctx, fmt.Sprintf("pick files response data type error: %T", respData))
		return nil
	}

	var result []string
	lo.ForEach(respData.([]any), func(file any, _ int) {
		result = append(result, file.(string))
	})
	return result
}

func getShowAppParams(ctx context.Context, selectAll bool) map[string]any {
	return map[string]any{
		"SelectAll":      selectAll,
		"Position":       NewMouseScreenPosition(),
		"QueryHistories": setting.GetSettingManager().GetLatestQueryHistory(ctx, 10),
		"LastQueryMode":  setting.GetSettingManager().GetWoxSetting(ctx).LastQueryMode,
	}
}

func onUIRequest(ctx context.Context, request WebsocketMsg) {
	if request.Method != "Log" {
		logger.Debug(ctx, fmt.Sprintf("got <%s> request from ui", request.Method))
	}

	switch request.Method {
	case "Ping":
		responseUISuccessWithData(ctx, request, "Pong")
	case "Log":
		handleUILog(ctx, request)
	case "Query":
		handleQuery(ctx, request)
	case "Action":
		handleAction(ctx, request)
	case "Refresh":
		handleRefresh(ctx, request)
	case "ChangeLanguage":
		handleChangeLanguage(ctx, request)
	case "GetLanguageJson":
		handleGetLanguageJson(ctx, request)
	case "VisibilityChanged":
		handleOnVisibilityChanged(ctx, request)
	case "LostFocus":
		handleLostFocus(ctx, request)
	case "GetQueryHistories":
		handleGetQueryHistories(ctx, request)
	}
}

func onUIResponse(ctx context.Context, response WebsocketMsg) {
	logger.Debug(ctx, fmt.Sprintf("got <%s> response from ui", response.Method))

	requestID := response.RequestId
	if requestID == "" {
		logger.Error(ctx, "response id not found")
		return
	}

	resultChan, exist := GetUIManager().GetUI(ctx).(*uiImpl).requestMap.Load(requestID)
	if !exist {
		logger.Error(ctx, fmt.Sprintf("response id not found: %s", requestID))
		return
	}

	resultChan <- response
}

func handleUILog(ctx context.Context, request WebsocketMsg) {
	traceId, traceIdErr := getWebsocketMsgParameter(ctx, request, "traceId")
	if traceIdErr != nil {
		logger.Error(ctx, traceIdErr.Error())
		responseUIError(ctx, request, traceIdErr.Error())
		return
	}
	level, levelErr := getWebsocketMsgParameter(ctx, request, "level")
	if levelErr != nil {
		logger.Error(ctx, levelErr.Error())
		responseUIError(ctx, request, levelErr.Error())
		return
	}
	message, messageErr := getWebsocketMsgParameter(ctx, request, "message")
	if messageErr != nil {
		logger.Error(ctx, messageErr.Error())
		responseUIError(ctx, request, messageErr.Error())
		return
	}

	logCtx := util.NewComponentContext(util.NewTraceContextWith(traceId), " UI")

	switch level {
	case "debug":
		logger.Debug(logCtx, message)
	case "info":
		logger.Info(logCtx, message)
	case "warn":
		logger.Warn(logCtx, message)
	case "error":
		logger.Error(logCtx, message)
	default:
		logger.Error(ctx, fmt.Sprintf("unsupported log level: %s", level))
		responseUIError(ctx, request, fmt.Sprintf("unsupported log level: %s", level))
	}
	responseUISuccess(ctx, request)
}

func handleQuery(ctx context.Context, request WebsocketMsg) {
	queryId, queryIdErr := getWebsocketMsgParameter(ctx, request, "queryId")
	if queryIdErr != nil {
		logger.Error(ctx, queryIdErr.Error())
		responseUIError(ctx, request, queryIdErr.Error())
		return
	}
	queryType, queryTypeErr := getWebsocketMsgParameter(ctx, request, "queryType")
	if queryTypeErr != nil {
		logger.Error(ctx, queryTypeErr.Error())
		responseUIError(ctx, request, queryTypeErr.Error())
		return
	}
	queryText, queryTextErr := getWebsocketMsgParameter(ctx, request, "queryText")
	if queryTextErr != nil {
		logger.Error(ctx, queryTextErr.Error())
		responseUIError(ctx, request, queryTextErr.Error())
		return
	}
	querySelectionJson, querySelectionErr := getWebsocketMsgParameter(ctx, request, "querySelection")
	if querySelectionErr != nil {
		logger.Error(ctx, querySelectionErr.Error())
		responseUIError(ctx, request, querySelectionErr.Error())
		return
	}
	var querySelection util.Selection
	json.Unmarshal([]byte(querySelectionJson), &querySelection)

	var changedQuery share.ChangedQuery
	if queryType == plugin.QueryTypeInput {
		changedQuery = share.ChangedQuery{
			QueryType: plugin.QueryTypeInput,
			QueryText: queryText,
		}
	} else if queryType == plugin.QueryTypeSelection {
		changedQuery = share.ChangedQuery{
			QueryType:      plugin.QueryTypeSelection,
			QueryText:      queryText,
			QuerySelection: querySelection,
		}
	} else {
		logger.Error(ctx, fmt.Sprintf("unsupported query type: %s", queryType))
		responseUIError(ctx, request, fmt.Sprintf("unsupported query type: %s", queryType))
		return
	}

	logger.Info(ctx, fmt.Sprintf("start to handle query changed: %s, queryId: %s", changedQuery.String(), queryId))

	if changedQuery.QueryType == plugin.QueryTypeInput && changedQuery.QueryText == "" {
		responseUISuccessWithData(ctx, request, []string{})
		return
	}
	if changedQuery.QueryType == plugin.QueryTypeSelection && changedQuery.QuerySelection.String() == "" {
		responseUISuccessWithData(ctx, request, []string{})
		return
	}

	query, queryErr := plugin.GetPluginManager().NewQuery(ctx, changedQuery)
	if queryErr != nil {
		logger.Error(ctx, queryErr.Error())
		responseUIError(ctx, request, queryErr.Error())
		return
	}

	var totalResultCount int
	var startTimestamp = util.GetSystemTimestamp()
	var resultDebouncer = util.NewDebouncer(24, func(results []plugin.QueryResultUI, reason string) {
		logger.Info(ctx, fmt.Sprintf("query %s: %s, result flushed (reason: %s), total results: %d", query.Type, query.String(), reason, totalResultCount))
		responseUISuccessWithData(ctx, request, results)
	})
	resultDebouncer.Start(ctx)
	logger.Info(ctx, fmt.Sprintf("query %s: %s, result flushed (new start)", query.Type, query.String()))
	resultChan, doneChan := plugin.GetPluginManager().Query(ctx, query)
	for {
		select {
		case results := <-resultChan:
			if len(results) == 0 {
				continue
			}
			lo.ForEach(results, func(_ plugin.QueryResultUI, index int) {
				results[index].QueryId = queryId
			})
			totalResultCount += len(results)
			resultDebouncer.Add(ctx, results)
		case <-doneChan:
			logger.Info(ctx, fmt.Sprintf("query done, total results: %d, cost %d ms", totalResultCount, util.GetSystemTimestamp()-startTimestamp))

			// if there is no result, show fallback websearch
			if totalResultCount == 0 {
				fallbackResults := plugin.GetPluginManager().QueryFallback(ctx, query)
				if len(fallbackResults) > 0 {
					lo.ForEach(fallbackResults, func(_ plugin.QueryResultUI, index int) {
						fallbackResults[index].QueryId = queryId
					})
					resultDebouncer.Add(ctx, fallbackResults)
					logger.Info(ctx, fmt.Sprintf("no result, show %d fallback results", len(fallbackResults)))
				} else {
					logger.Info(ctx, "no result, no fallback results")
				}
			}

			resultDebouncer.Done(ctx)
			return
		case <-time.After(time.Minute):
			logger.Info(ctx, fmt.Sprintf("query timeout, query: %s, request id: %s", query, request.RequestId))
			resultDebouncer.Done(ctx)
			responseUIError(ctx, request, fmt.Sprintf("query timeout, query: %s, request id: %s", query, request.RequestId))
			return
		}
	}

}

func handleAction(ctx context.Context, request WebsocketMsg) {
	resultId, idErr := getWebsocketMsgParameter(ctx, request, "resultId")
	if idErr != nil {
		logger.Error(ctx, idErr.Error())
		responseUIError(ctx, request, idErr.Error())
		return
	}
	actionId, actionIdErr := getWebsocketMsgParameter(ctx, request, "actionId")
	if actionIdErr != nil {
		logger.Error(ctx, actionIdErr.Error())
		responseUIError(ctx, request, actionIdErr.Error())
		return
	}

	plugin.GetPluginManager().ExecuteAction(ctx, resultId, actionId)

	responseUISuccess(ctx, request)
}

func handleRefresh(ctx context.Context, request WebsocketMsg) {
	resultStr, resultErr := getWebsocketMsgParameter(ctx, request, "refreshableResult")
	if resultErr != nil {
		logger.Error(ctx, resultErr.Error())
		responseUIError(ctx, request, resultErr.Error())
		return
	}

	queryId, queryIdErr := getWebsocketMsgParameter(ctx, request, "queryId")
	if queryIdErr != nil {
		logger.Error(ctx, queryIdErr.Error())
		responseUIError(ctx, request, queryIdErr.Error())
		return
	}

	var result plugin.RefreshableResultWithResultId
	unmarshalErr := json.Unmarshal([]byte(resultStr), &result)
	if unmarshalErr != nil {
		logger.Error(ctx, unmarshalErr.Error())
		responseUIError(ctx, request, unmarshalErr.Error())
		return
	}

	startTime := util.GetSystemTimestamp()
	logger.Debug(ctx, fmt.Sprintf("start executing refresh for result: %s (resultId:%s, queryId:%s)", result.Title, result.ResultId, queryId))

	// replace remote preview with local preview
	if result.Preview.PreviewType == plugin.WoxPreviewTypeRemote {
		preview, err := plugin.GetPluginManager().GetResultPreview(util.NewTraceContext(), result.ResultId)
		if err != nil {
			logger.Error(ctx, err.Error())
			responseUIError(ctx, request, err.Error())
			return
		}
		result.Preview = preview
	}

	newResult, refreshErr := plugin.GetPluginManager().ExecuteRefresh(ctx, result)
	logger.Debug(ctx, fmt.Sprintf("finished refresh %s, cost: %dms", result.ResultId, util.GetSystemTimestamp()-startTime))
	if refreshErr != nil {
		logger.Error(ctx, refreshErr.Error())
		responseUIError(ctx, request, refreshErr.Error())
		return
	}

	responseUISuccessWithData(ctx, request, newResult)
}

func handleChangeLanguage(ctx context.Context, request WebsocketMsg) {
	langCode, langCodeErr := getWebsocketMsgParameter(ctx, request, "langCode")
	if langCodeErr != nil {
		logger.Error(ctx, langCodeErr.Error())
		responseUIError(ctx, request, langCodeErr.Error())
		return
	}
	if !i18n.IsSupportedLangCode(langCode) {
		logger.Error(ctx, fmt.Sprintf("unsupported lang code: %s", langCode))
		responseUIError(ctx, request, fmt.Sprintf("unsupported lang code: %s", langCode))
		return
	}

	langErr := i18n.GetI18nManager().UpdateLang(ctx, i18n.LangCode(langCode))
	if langErr != nil {
		logger.Error(ctx, langErr.Error())
		responseUIError(ctx, request, langErr.Error())
		return
	}

	responseUISuccess(ctx, request)
}

func handleGetLanguageJson(ctx context.Context, request WebsocketMsg) {
	langCode, langCodeErr := getWebsocketMsgParameter(ctx, request, "langCode")
	if langCodeErr != nil {
		logger.Error(ctx, langCodeErr.Error())
		responseUIError(ctx, request, langCodeErr.Error())
		return
	}
	if !i18n.IsSupportedLangCode(langCode) {
		logger.Error(ctx, fmt.Sprintf("unsupported lang code: %s", langCode))
		responseUIError(ctx, request, fmt.Sprintf("unsupported lang code: %s", langCode))
		return
	}

	langJson, err := i18n.GetI18nManager().GetLangJson(ctx, i18n.LangCode(langCode))
	if err != nil {
		logger.Error(ctx, err.Error())
		responseUIError(ctx, request, err.Error())
		return
	}

	responseUISuccessWithData(ctx, request, langJson)
}

func handleOnVisibilityChanged(ctx context.Context, request WebsocketMsg) {
	isVisible, isVisibleErr := getWebsocketMsgParameter(ctx, request, "isVisible")
	if isVisibleErr != nil {
		logger.Error(ctx, isVisibleErr.Error())
		responseUIError(ctx, request, isVisibleErr.Error())
		return
	}

	query, queryErr := getWebsocketMsgParameter(ctx, request, "query")
	if queryErr != nil {
		logger.Error(ctx, queryErr.Error())
		responseUIError(ctx, request, queryErr.Error())
		return
	}
	var changedQuery share.ChangedQuery
	unmarshalErr := json.Unmarshal([]byte(query), &changedQuery)
	if unmarshalErr != nil {
		logger.Error(ctx, unmarshalErr.Error())
		responseUIError(ctx, request, unmarshalErr.Error())
		return
	}

	if isVisible == "true" {
		onAppShow(ctx)
	} else {
		onAppHide(ctx, changedQuery)
	}
}

func handleLostFocus(ctx context.Context, request WebsocketMsg) {
	woxSetting := setting.GetSettingManager().GetWoxSetting(ctx)
	if woxSetting.HideOnLostFocus {
		GetUIManager().GetUI(ctx).HideApp(ctx)
	}
}

func handleGetQueryHistories(ctx context.Context, request WebsocketMsg) {
	queryHistories := setting.GetSettingManager().GetWoxAppData(ctx).QueryHistories
	responseUISuccessWithData(ctx, request, queryHistories)
}

func onAppShow(ctx context.Context) {
	woxSetting := setting.GetSettingManager().GetWoxSetting(ctx)
	if woxSetting.SwitchInputMethodABC {
		util.GetLogger().Info(ctx, "switch input method to ABC")
		ime.SwitchInputMethodABC()
	}
}

func onAppHide(ctx context.Context, query share.ChangedQuery) {
	setting.GetSettingManager().AddQueryHistory(ctx, query)
	if setting.GetSettingManager().GetWoxSetting(ctx).LastQueryMode == setting.LastQueryModeEmpty {
		GetUIManager().GetUI(ctx).ChangeQuery(ctx, share.ChangedQuery{
			QueryType: plugin.QueryTypeInput,
			QueryText: "",
		})
	}
}

func getWebsocketMsgParameter(ctx context.Context, msg WebsocketMsg, key string) (string, error) {
	jsonData, marshalErr := json.Marshal(msg.Data)
	if marshalErr != nil {
		return "", marshalErr
	}

	paramterData := gjson.GetBytes(jsonData, key)
	if !paramterData.Exists() {
		return "", errors.New(fmt.Sprintf("%s parameter not found", key))
	}

	return paramterData.String(), nil
}

func getWindowMouseScreenLocation(windowWidth int) (int, int) {
	size := screen.GetMouseScreen()
	x := size.X + (size.Width-windowWidth)/2
	y := size.Height / 6
	return x, y
}