Auth Files and Components

"use client";

import { type ReactNode, useEffect, useState } from "react";
import { haWebSocket } from "@/lib/haWebsocket";
import { getAccessToken, login, exchangeCodeForToken, ENV } from "@/lib/haAuth";

export function HomeAssistantProvider({ children, useProxy = false }: { children: ReactNode; useProxy?: boolean }) {
    const [ready, setReady] = useState(false);
    const [token, setToken] = useState<string | null>(() => getAccessToken());

    useEffect(() => {
        if (token) return;

        const searchParams = new URLSearchParams(window.location.search);
        const code = searchParams.get("code");

        if (code) {
            exchangeCodeForToken(code)
                .then(() => {
                    setToken(getAccessToken());
                    // Clean Code from URL
                    window.history.replaceState({}, "", window.location.pathname);
                })
                .catch((err) => {
                    console.error("Failed to exchange code for token:", err);
                });
        } else {
            login(); // redirect to HA OAuth2
        }
    }, []);

    useEffect(() => {
        if (!token) return;
        if (useProxy && !ENV.PROXY_URL) {
            console.error("PROXY_URL is required when useProxy is enabled");
            return;
        }

        async function connect() {
            try {
                await haWebSocket.connect(
                    useProxy ? ENV.PROXY_URL! : `${ENV.HA_HOST}:${ENV.HA_PORT}`,
                    token!,
                    useProxy,
                );
                setReady(true);
            } catch (err) {
                console.error("Failed to connect to HA WebSocket:", err);
            }
        }

        connect();
    }, [token, useProxy]);

    if (!ready) {
        return <div>Connecting to Home Assistant...</div>;
    }

    return <>{children}</>;
}

/**
 * Home Assistant Authentication Utility
 *
 * Supports OAuth2 token exchange, refresh, persistent storage in localStorage,
 * and an optional long-lived token override (NEXT_PUBLIC_HA_LONG_LIVED_TOKEN).
 */

export const ENV = {
    // @ts-ignore - import.meta may not exist in Next
    HA_HOST: typeof process !== "undefined" ? process.env.NEXT_PUBLIC_HA_URL : import.meta.env.VITE_HA_URL, // e.g. homeassistant.local

    // @ts-ignore - import.meta may not exist in Next
    HA_PORT: typeof process !== "undefined" ? process.env.NEXT_PUBLIC_HA_PORT : import.meta.env.VITE_HA_PORT, // e.g. 8123

    // OAuth2 app credentials (must match Home Assistant config)
    CLIENT_ID:
        // @ts-ignore - import.meta may not exist in Next
        typeof process !== "undefined" ? process.env.NEXT_PUBLIC_HA_CLIENT_ID : import.meta.env.VITE_HA_CLIENT_ID,

    // Optional override: if set, skip OAuth2
    LONG_LIVED_TOKEN:
        typeof process !== "undefined"
            ? process.env.NEXT_PUBLIC_HA_LONG_LIVED_TOKEN
            : // @ts-ignore - import.meta may not exist in Next
              import.meta.env.VITE_HA_LONG_LIVED_TOKEN,

    PROXY_URL:
        typeof process !== "undefined"
            ? process.env.NEXT_PUBLIC_PROXY_URL
            : // @ts-ignore - import.meta may not exist in Next
              import.meta.env.VITE_PROXY_URL,
};

// Validate required environment variables at module load
const requiredEnvVars = [
    { key: "HA_HOST", value: ENV.HA_HOST },
    { key: "HA_PORT", value: ENV.HA_PORT },
    { key: "CLIENT_ID", value: ENV.CLIENT_ID },
];

const missingVars = requiredEnvVars.filter(({ value }) => !value);
if (missingVars.length > 0) {
    const missing = missingVars.map(({ key }) => key).join(", ");
    throw new Error(
        `Missing required environment variables: ${missing}. ` +
            `Please set NEXT_PUBLIC_* or VITE_* prefixed versions.`,
    );
}

const HA_HTTP_URL = `http://${ENV.HA_HOST}:${ENV.HA_PORT}`;

const STORAGE_KEY = "ha_tokens";

export type Tokens = {
    access_token: string;
    refresh_token?: string;
    expires_in?: number; // in seconds
    token_type: string;
    created_at: number; // epoch ms
};

function log(...args: any[]) {
    console.log("[HA AUTH]", ...args);
}

/**
 * Save tokens into localStorage (with created_at timestamp).
 */
export function saveTokens(tokens: any) {
    const data: Tokens = {
        ...tokens,
        created_at: Date.now(),
    };
    localStorage.setItem(STORAGE_KEY, JSON.stringify(data));
    log("Saved tokens to localStorage:", {
        access_token: !!data.access_token,
        refresh_token: !!data.refresh_token,
        expires_in: data.expires_in,
    });
}

/**
 * Load tokens from localStorage.
 */
export function loadTokens(): Tokens | null {
    try {
        const raw = localStorage.getItem(STORAGE_KEY);
        return raw ? (JSON.parse(raw) as Tokens) : null;
    } catch {
        return null;
    }
}

/**
 * Clear stored tokens.
 */
export function clearTokens() {
    localStorage.removeItem(STORAGE_KEY);
    log("Cleared tokens");
}

/**
 * Redirect to Home Assistant OAuth2 login.
 */
export function login() {
    if (ENV.LONG_LIVED_TOKEN) {
        log("Using long-lived token; skipping OAuth login");
        return;
    }
    const authUrl =
        `${HA_HTTP_URL}/auth/authorize?response_type=code` +
        `&client_id=${encodeURIComponent(ENV.CLIENT_ID)}` +
        `&redirect_uri=${encodeURIComponent(window.location.href)}`;

    log("Redirecting to HA login:", authUrl);

    window.location.href = authUrl;
}

/**
 * Exchange authorization code for access + refresh tokens.
 */
export async function exchangeCodeForToken(code: string) {
    log("Exchanging code for token");
    const res = await fetch(`${HA_HTTP_URL}/auth/token`, {
        method: "POST",
        headers: { "Content-Type": "application/x-www-form-urlencoded" },
        body: new URLSearchParams({
            grant_type: "authorization_code",
            code,
            client_id: ENV.CLIENT_ID,
            redirect_uri: window.location.href,
        }),
    });

    const text = await res.text();
    if (!res.ok) {
        log("Token exchange failed:", res.status, text);
        throw new Error(`Token exchange failed: ${res.status}`);
    }
    const data = JSON.parse(text);
    saveTokens(data);
    return data as Tokens;
}

/**
 * Refresh access token using stored refresh token.
 */
export async function refreshAccessToken() {
    const tokens = loadTokens();
    if (!tokens?.refresh_token) throw new Error("No refresh token available");

    log("Refreshing access token...");
    const res = await fetch(`${HA_HTTP_URL}/auth/token`, {
        method: "POST",
        headers: { "Content-Type": "application/x-www-form-urlencoded" },
        body: new URLSearchParams({
            grant_type: "refresh_token",
            refresh_token: tokens.refresh_token,
            client_id: ENV.CLIENT_ID,
        }),
    });

    const text = await res.text();
    if (!res.ok) {
        log("Refresh failed:", res.status, text);
        throw new Error(`Refresh failed: ${res.status}`);
    }
    const data = JSON.parse(text);
    saveTokens(data);
    log("Refresh successful; new access token saved");
    return data as Tokens;
}

/**
 * Get current access token (returns null if expired or missing).
 */
export function getAccessToken(): string | null {
    if (ENV.LONG_LIVED_TOKEN) return ENV.LONG_LIVED_TOKEN;

    const t = loadTokens();
    if (!t) return null;

    if (!t.expires_in) return t.access_token;

    const expiresAt = t.created_at + t.expires_in * 1000;
    if (Date.now() >= expiresAt) {
        log("Access token expired");
        return null;
    }
    return t.access_token;
}

/**
 * Get time (ms) until current access token expires.
 */
export function getTimeToExpiryMs(): number | null {
    const t = loadTokens();
    if (!t?.expires_in) return null;
    return t.created_at + t.expires_in * 1000 - Date.now();
}

/**
 * Force a refresh now (useful for debugging).
 */
export async function forceRefreshNow() {
    return refreshAccessToken();
}

import { refreshAccessToken, clearTokens, login, getTimeToExpiryMs, ENV } from "@/lib/haAuth";

class HAWebSocket {
    private ws: WebSocket | null = null;
    private msgId = 1;
    private refreshTimer: number | null = null;

    private stateListeners: Map<string, ((state: any) => void)[]> = new Map();

    async connect(baseUrl: string, token: string, useProxy: boolean = false): Promise<void> {
        return new Promise((resolve, reject) => {
            const HA_HOST = ENV.HA_HOST;
            const HA_PORT = ENV.HA_PORT;
            const haWsUrl = `ws://${HA_HOST}:${HA_PORT}/api/websocket`;
            const fullUrl = useProxy ? `${baseUrl}/ha` : haWsUrl;

            console.log("[HA WS] Connecting:", fullUrl);
            this.ws = new WebSocket(fullUrl);

            this.ws.onopen = () => console.log("[HA WS] Opened");

            this.ws.onmessage = async (event) => {
                const msg = JSON.parse(event.data);

                if (msg.type === "auth_required") {
                    this.ws?.send(JSON.stringify({ type: "auth", access_token: token }));
                } else if (msg.type === "auth_ok") {
                    console.log("[HA WS] Authenticated");
                    this.subscribeEvents("state_changed");
                    this.scheduleRefresh(baseUrl, useProxy);
                    resolve();
                } else if (msg.type === "auth_invalid") {
                    console.warn("[HA WS] auth_invalid, trying refresh");
                    try {
                        const newTokens = await refreshAccessToken();
                        const newToken = newTokens.access_token;
                        console.log("[HA WS] Token refreshed, retrying connect");
                        await this.reconnect(baseUrl, newToken, useProxy);
                        resolve();
                    } catch (err) {
                        console.error("[HA WS] Refresh failed; clearing tokens and restarting login");
                        clearTokens();
                        login();
                        reject(err);
                    }
                } else if (msg.type === "event" && msg.event?.event_type === "state_changed") {
                    const entityId = msg.event.data.entity_id;
                    this.notifyStateListeners(entityId, msg.event.data.new_state);
                }
            };

            this.ws.onerror = (err) => {
                console.error("[HA WS] WebSocket error:", err);
                reject(err);
            };

            this.ws.onclose = () => {
                console.warn("[HA WS] Closed");
                this.clearTimers();
            };
        });
    }

    private async reconnect(baseUrl: string, token: string, useProxy: boolean) {
        this.clearTimers();
        try {
            if (this.ws && this.ws.readyState === WebSocket.OPEN) this.ws.close();
        } catch {}
        return this.connect(baseUrl, token, useProxy);
    }

    private scheduleRefresh(baseUrl: string, useProxy: boolean) {
        this.clearRefreshTimer();
        const tte = getTimeToExpiryMs();
        if (tte == null) {
            console.log("[HA WS] Token has no expiry; skipping proactive refresh (likely long-lived)");
            return;
        }
        const refreshIn = Math.max(tte - 5 * 60 * 1000, 10 * 1000); // 5 min before expiry, at least 10s
        console.log("[HA WS] Scheduling refresh in", Math.round(refreshIn / 1000), "s");
        this.refreshTimer = window.setTimeout(async () => {
            try {
                const newTokens = await refreshAccessToken();
                const newToken = newTokens.access_token;
                console.log("[HA WS] Proactive refresh OK; reconnecting WS");
                await this.reconnect(baseUrl, newToken, useProxy);
            } catch (e) {
                console.warn("[HA WS] Proactive refresh failed; restarting auth");
                clearTokens();
                login();
            }
        }, refreshIn);
    }

    private subscribeEvents(eventType: string) {
        const id = ++this.msgId;
        this.ws?.send(
            JSON.stringify({
                id,
                type: "subscribe_events",
                event_type: eventType,
            }),
        );
    }

    private clearRefreshTimer() {
        if (this.refreshTimer) window.clearTimeout(this.refreshTimer);
        this.refreshTimer = null;
    }

    private clearTimers() {
        this.clearRefreshTimer();
    }

    getState(entityId: string): Promise<any> {
        return new Promise((resolve, reject) => {
            const id = ++this.msgId;
            const TIMEOUT_MS = 10000;

            // Reject immediately if WebSocket is not connected
            if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
                reject(new Error("WebSocket is not connected"));
                return;
            }

            // Set timeout to prevent hanging
            const timeoutId = setTimeout(() => {
                this.ws?.removeEventListener("message", handleMessage);
                reject(new Error(`getState timeout for entity: ${entityId}`));
            }, TIMEOUT_MS);

            const handleMessage = (event: MessageEvent) => {
                const msg = JSON.parse(event.data);
                if (msg.id === id && msg.type === "result") {
                    clearTimeout(timeoutId);
                    this.ws?.removeEventListener("message", handleMessage);
                    if (msg.success && Array.isArray(msg.result)) {
                        const entity = msg.result.find((e: any) => e.entity_id === entityId);
                        resolve(entity);
                    } else {
                        resolve(null);
                    }
                }
            };

            this.ws?.addEventListener("message", handleMessage);

            this.ws?.send(
                JSON.stringify({
                    id,
                    type: "get_states",
                }),
            );
        });
    }

    addStateListener(entityId: string, cb: (newState: any) => void) {
        const listeners = this.stateListeners.get(entityId) || [];
        listeners.push(cb);
        this.stateListeners.set(entityId, listeners);
    }

    removeStateListener(entityId: string, cb: (newState: any) => void) {
        let listeners = this.stateListeners.get(entityId) || [];
        listeners = listeners.filter((fn) => fn !== cb);
        this.stateListeners.set(entityId, listeners);
    }

    private notifyStateListeners(entityId: string, newState: any) {
        const listeners = this.stateListeners.get(entityId);
        if (listeners) listeners.forEach((fn) => fn(newState));
    }

    callService(domain: string, service: string, serviceData: any) {
        const id = ++this.msgId;
        this.ws?.send(
            JSON.stringify({
                id,
                type: "call_service",
                domain,
                service,
                service_data: serviceData,
            }),
        );
    }
    callServiceWithResponse(domain: string, service: string, serviceData: Record<string, any> = {}): Promise<any> {
        return new Promise((resolve, reject) => {
            const id = ++this.msgId;

            const handleMessage = (event: MessageEvent) => {
                const msg = JSON.parse(event.data);
                if (msg.id === id && msg.type === "result") {
                    this.ws?.removeEventListener("message", handleMessage);
                    if (msg.success) {
                        // The response will be in msg.result
                        resolve(msg.result);
                    } else {
                        reject(msg.error || new Error("Service call failed"));
                    }
                }
            };

            this.ws?.addEventListener("message", handleMessage);

            // According to HA API, set `return_response: true` to get a payload back.
            this.ws?.send(
                JSON.stringify({
                    id,
                    type: "call_service",
                    domain,
                    service,
                    service_data: serviceData,
                    return_response: true,
                }),
            );
        });
    }
    fireEvent(eventType: string, eventData: Record<string, any> = {}): Promise<void> {
        return new Promise((resolve, reject) => {
            const id = ++this.msgId;

            const handleMessage = (event: MessageEvent) => {
                const msg = JSON.parse(event.data);
                if (msg.id === id && msg.type === "result") {
                    this.ws?.removeEventListener("message", handleMessage);
                    if (msg.success) {
                        resolve();
                    } else {
                        reject(msg.error);
                    }
                }
            };

            this.ws?.addEventListener("message", handleMessage);

            this.ws?.send(
                JSON.stringify({
                    id,
                    type: "fire_event",
                    event_type: eventType,
                    event_data: eventData,
                }),
            );
        });
    }
}

export const haWebSocket = new HAWebSocket();