EMS_Platform/backend/utils/ha_client.go

package utils

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strings"
	"time"

	"gorm.io/datatypes"
)

type HAConfig struct {
	URL   string `json:"url"`
	Token string `json:"token"`
}

type HAEntity struct {
	EntityID    string                 `json:"entity_id"`
	State       string                 `json:"state"`
	Attributes  map[string]interface{} `json:"attributes"`
	LastChanged time.Time              `json:"last_changed"`
	LastUpdated time.Time              `json:"last_updated"`
	DeviceID    string                 `json:"device_id"`
	DeviceName  string                 `json:"device_name"`
}

type HATemplateReq struct {
	Template string `json:"template"`
}

type HATemplateResult struct {
	ID    string `json:"id"`
	State string `json:"s"`
	Name  string `json:"n"`
	DID   string `json:"did"`
	DName string `json:"dn"`
}

// FetchHAEntitiesByDevice fetches entities for a specific device using HA Template API
func FetchHAEntitiesByDevice(config datatypes.JSON, deviceID string) ([]HAEntity, error) {
	haConfig, err := parseHAConfig(config)
	if err != nil {
		return nil, err
	}

	client := &http.Client{Timeout: 10 * time.Second}
	url := normalizeURL(haConfig.URL)

	rawTemplate := `
{% set ns = namespace(result=[]) %}
{% set device_entities = device_entities('__DEVICE_ID__') %}
{% for entity_id in device_entities %}
  {% set state = states[entity_id] %}
  {% if state %}
    {% set name = state.attributes.friendly_name %}
    {% if name is not defined or name is none %}
      {% set name = entity_id %}
    {% endif %}
    {% set entry = {
      "id": entity_id,
      "s": state.state,
      "n": name,
      "did": '__DEVICE_ID__',
      "dn": ''
    } %}
    {% set ns.result = ns.result + [entry] %}
  {% endif %}
{% endfor %}
{{ ns.result | to_json }}
`
	template := strings.ReplaceAll(rawTemplate, "__DEVICE_ID__", deviceID)

	return executeTemplateQuery(client, url, haConfig.Token, template)
}

// FetchEntityState fetches a single entity state
func FetchEntityState(config datatypes.JSON, entityID string) (*HAEntity, error) {
	haConfig, err := parseHAConfig(config)
	if err != nil {
		return nil, err
	}

	client := &http.Client{Timeout: 5 * time.Second}
	url := normalizeURL(haConfig.URL)
	
	req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/states/%s", url, entityID), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+haConfig.Token)
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("HA returned status: %s", resp.Status)
	}

	var entity HAEntity
	if err := json.NewDecoder(resp.Body).Decode(&entity); err != nil {
		return nil, err
	}
	return &entity, nil
}

// BatchFetchStates fetches multiple entities in one go using template
func BatchFetchStates(config datatypes.JSON, entityIDs []string) (map[string]string, error) {
	if len(entityIDs) == 0 {
		return map[string]string{}, nil
	}
	
	haConfig, err := parseHAConfig(config)
	if err != nil {
		return nil, err
	}

	client := &http.Client{Timeout: 10 * time.Second}
	url := normalizeURL(haConfig.URL)

	// Safer template approach: Return List of Structs instead of Dict
	// This avoids potential Jinja2 dictionary key issues or namespace update quirks
	idsJson, _ := json.Marshal(entityIDs)
	// Escape % for Sprintf by using %%
	template := fmt.Sprintf(`
{%% set ids = %s %%}
{%% set result = [] %%}
{%% for id in ids %%}
  {%% set s = states(id) %%}
  {%% if s not in ['unknown', 'unavailable', 'none', ''] %%}
    {%% set result = result + [{"id": id, "s": s}] %%}
  {%% endif %%}
{%% endfor %%}
{{ result | to_json }}
`, string(idsJson))

	reqBody, _ := json.Marshal(HATemplateReq{Template: template})
	req, err := http.NewRequest("POST", url+"/api/template", bytes.NewBuffer(reqBody))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+haConfig.Token)
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Read body for debug if error
	if resp.StatusCode != 200 {
		bodyBytes, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("HA error: %s, Body: %s", resp.Status, string(bodyBytes))
	}

	// Helper struct for decoding the list response
	type StateResult struct {
		ID    string `json:"id"`
		State string `json:"s"`
	}

	var results []StateResult
	if err := json.NewDecoder(resp.Body).Decode(&results); err != nil {
		return nil, fmt.Errorf("decode error: %v", err)
	}

	// Convert list to map
	resultMap := make(map[string]string)
	for _, r := range results {
		resultMap[r.ID] = r.State
	}

	return resultMap, nil
}

// Helpers

func parseHAConfig(config datatypes.JSON) (HAConfig, error) {
	var haConfig HAConfig
	b, err := config.MarshalJSON()
	if err != nil {
		return haConfig, fmt.Errorf("config error: %v", err)
	}
	if err := json.Unmarshal(b, &haConfig); err != nil {
		return haConfig, fmt.Errorf("invalid configuration format: %v", err)
	}
	if haConfig.URL == "" || haConfig.Token == "" {
		return haConfig, fmt.Errorf("URL and Token are required")
	}
	return haConfig, nil
}

func normalizeURL(url string) string {
	url = strings.TrimSuffix(url, "/")
	url = strings.TrimSuffix(url, "/api")
	return url
}

func executeTemplateQuery(client *http.Client, url, token, template string) ([]HAEntity, error) {
	reqBody, _ := json.Marshal(HATemplateReq{Template: template})
	req, err := http.NewRequest("POST", url+"/api/template", bytes.NewBuffer(reqBody))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %v", err)
	}
	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("connection failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Home Assistant returned status: %s", resp.Status)
	}

	var tmplResults []HATemplateResult
	if err := json.NewDecoder(resp.Body).Decode(&tmplResults); err != nil {
		return nil, fmt.Errorf("failed to decode response: %v", err)
	}

	entities := make([]HAEntity, len(tmplResults))
	for i, r := range tmplResults {
		entities[i] = HAEntity{
			EntityID:   r.ID,
			State:      r.State,
			Attributes: map[string]interface{}{"friendly_name": r.Name},
			DeviceID:   r.DID,
			DeviceName: r.DName,
		}
	}

	sort.Slice(entities, func(i, j int) bool {
		nameI, _ := entities[i].Attributes["friendly_name"].(string)
		nameJ, _ := entities[j].Attributes["friendly_name"].(string)
		if nameI == "" { nameI = entities[i].EntityID }
		if nameJ == "" { nameJ = entities[j].EntityID }
		return nameI < nameJ
	})

	return entities, nil
}