// Hermes chat tenant router assigns each Keycloak subject to one isolated
// Hermes pod and exposes only the user-facing parts of Hermes WebUI.
package main

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base32"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

type linkRecord struct {
	Slot      int   `json:"slot"`
	ExpiresAt int64 `json:"expires_at"`
}

type tenantState struct {
	Salt           string                `json:"salt"`
	Assignments    map[string]int        `json:"assignments"`
	Telegram       map[string]int        `json:"telegram,omitempty"`
	LinkCodes      map[string]linkRecord `json:"link_codes,omitempty"`
	TelegramOffset int64                 `json:"telegram_offset,omitempty"`
}

type tenantRouter struct {
	mu            sync.Mutex
	state         tenantState
	statePath     string
	slots         int
	backendURL    func(int) string
	backendAPIURL func(int) string
	now           func() time.Time
	telegram      *telegramBot
}

var deniedPrefixes = []string{
	"/api/admin", "/api/commands/exec", "/api/config", "/api/console", "/api/credentials",
	"/api/cron", "/api/curator", "/api/dashboard/config", "/api/env",
	"/api/escape", "/api/extensions", "/api/file", "/api/folder",
	"/api/gateway", "/api/git", "/api/git-info", "/api/health/restart", "/api/kanban", "/api/logs", "/api/mcp",
	"/api/messaging", "/api/notes", "/api/ops",
	"/api/oauth", "/api/onboarding/oauth", "/api/pairing", "/api/plugins", "/api/providers",
	"/api/rollback", "/api/shutdown", "/api/terminal",
	"/api/tools", "/api/updates", "/api/webhooks", "/api/wiki",
	"/api/workspace", "/api/workspaces",
}

const retiredServiceWorker = `self.addEventListener("install",()=>self.skipWaiting());self.addEventListener("activate",event=>event.waitUntil((async()=>{for(const key of await caches.keys())await caches.delete(key);await self.registration.unregister();for(const client of await self.clients.matchAll({type:"window"}))await client.navigate(client.url)})()));`

const (
	trustedTenantHeader = "X-Hermes-Tenant-Identity"
	tenantSessionCookie = "hermes_chat_session"
	tenantProfileCookie = "hermes_chat_profile"
)

func newTenantRouter(statePath string, slots int, backendURL func(int) string) (*tenantRouter, error) {
	if slots < 1 {
		return nil, fmt.Errorf("tenant slots must be positive")
	}
	router := &tenantRouter{
		statePath:  statePath,
		slots:      slots,
		backendURL: backendURL,
		now:        time.Now,
		state: tenantState{
			Assignments: map[string]int{},
			Telegram:    map[string]int{},
			LinkCodes:   map[string]linkRecord{},
		},
	}
	content, err := os.ReadFile(statePath)
	if err == nil {
		if err := json.Unmarshal(content, &router.state); err != nil {
			return nil, fmt.Errorf("decode tenant state: %w", err)
		}
	} else if !os.IsNotExist(err) {
		return nil, fmt.Errorf("read tenant state: %w", err)
	}
	if router.state.Assignments == nil {
		router.state.Assignments = map[string]int{}
	}
	if router.state.Telegram == nil {
		router.state.Telegram = map[string]int{}
	}
	if router.state.LinkCodes == nil {
		router.state.LinkCodes = map[string]linkRecord{}
	}
	if router.state.Salt == "" {
		salt := make([]byte, 32)
		if _, err := rand.Read(salt); err != nil {
			return nil, fmt.Errorf("create tenant salt: %w", err)
		}
		router.state.Salt = hex.EncodeToString(salt)
		if err := router.saveLocked(); err != nil {
			return nil, err
		}
	}
	return router, nil
}

func (router *tenantRouter) saveLocked() error {
	content, err := json.MarshalIndent(router.state, "", "  ")
	if err != nil {
		return fmt.Errorf("encode tenant state: %w", err)
	}
	if err := os.MkdirAll(filepath.Dir(router.statePath), 0700); err != nil {
		return fmt.Errorf("create tenant state directory: %w", err)
	}
	temporary := router.statePath + ".tmp"
	if err := os.WriteFile(temporary, append(content, '\n'), 0600); err != nil {
		return fmt.Errorf("write tenant state: %w", err)
	}
	if err := os.Rename(temporary, router.statePath); err != nil {
		return fmt.Errorf("replace tenant state: %w", err)
	}
	return nil
}

func (router *tenantRouter) identityHash(kind, value string) string {
	digest := sha256.Sum256([]byte(router.state.Salt + "\x00" + kind + "\x00" + value))
	return hex.EncodeToString(digest[:])
}

func (router *tenantRouter) slotFor(subject string) (int, error) {
	router.mu.Lock()
	defer router.mu.Unlock()
	identity := router.identityHash("keycloak", subject)
	if slot, ok := router.state.Assignments[identity]; ok {
		return slot, nil
	}
	used := make(map[int]bool, len(router.state.Assignments))
	for _, slot := range router.state.Assignments {
		used[slot] = true
	}
	for slot := 0; slot < router.slots; slot++ {
		if used[slot] {
			continue
		}
		router.state.Assignments[identity] = slot
		if err := router.saveLocked(); err != nil {
			delete(router.state.Assignments, identity)
			return 0, err
		}
		return slot, nil
	}
	return 0, fmt.Errorf("all isolated chat slots are assigned")
}

func (router *tenantRouter) createLink(subject string) (string, time.Time, error) {
	slot, err := router.slotFor(subject)
	if err != nil {
		return "", time.Time{}, err
	}
	raw := make([]byte, 5)
	if _, err := rand.Read(raw); err != nil {
		return "", time.Time{}, fmt.Errorf("create link code: %w", err)
	}
	code := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(raw)
	expires := router.now().Add(10 * time.Minute)
	router.mu.Lock()
	defer router.mu.Unlock()
	for digest, record := range router.state.LinkCodes {
		if record.Slot == slot || record.ExpiresAt <= router.now().Unix() {
			delete(router.state.LinkCodes, digest)
		}
	}
	digest := router.identityHash("link", strings.ToUpper(code))
	router.state.LinkCodes[digest] = linkRecord{Slot: slot, ExpiresAt: expires.Unix()}
	if err := router.saveLocked(); err != nil {
		delete(router.state.LinkCodes, digest)
		return "", time.Time{}, err
	}
	return code, expires, nil
}

func (router *tenantRouter) consumeLink(telegramUser, code string) (int, error) {
	router.mu.Lock()
	defer router.mu.Unlock()
	digest := router.identityHash("link", strings.ToUpper(strings.TrimSpace(code)))
	record, ok := router.state.LinkCodes[digest]
	if !ok || record.ExpiresAt <= router.now().Unix() {
		delete(router.state.LinkCodes, digest)
		return 0, fmt.Errorf("link code is invalid or expired")
	}
	delete(router.state.LinkCodes, digest)
	router.state.Telegram[router.identityHash("telegram", telegramUser)] = record.Slot
	if err := router.saveLocked(); err != nil {
		return 0, err
	}
	return record.Slot, nil
}

func (router *tenantRouter) telegramSlot(telegramUser string) (int, bool) {
	router.mu.Lock()
	defer router.mu.Unlock()
	slot, ok := router.state.Telegram[router.identityHash("telegram", telegramUser)]
	return slot, ok
}

func (router *tenantRouter) telegramLinked(subject string) (bool, error) {
	slot, err := router.slotFor(subject)
	if err != nil {
		return false, err
	}
	router.mu.Lock()
	defer router.mu.Unlock()
	for _, linkedSlot := range router.state.Telegram {
		if linkedSlot == slot {
			return true, nil
		}
	}
	return false, nil
}

func (router *tenantRouter) unlinkTelegram(subject string) error {
	slot, err := router.slotFor(subject)
	if err != nil {
		return err
	}
	router.mu.Lock()
	defer router.mu.Unlock()
	for identity, linkedSlot := range router.state.Telegram {
		if linkedSlot == slot {
			delete(router.state.Telegram, identity)
		}
	}
	for digest, record := range router.state.LinkCodes {
		if record.Slot == slot {
			delete(router.state.LinkCodes, digest)
		}
	}
	return router.saveLocked()
}

func authenticatedSubject(request *http.Request) string {
	for _, header := range []string{"X-Forwarded-User", "X-Auth-Request-User"} {
		if subject := strings.TrimSpace(request.Header.Get(header)); subject != "" {
			return subject
		}
	}
	return ""
}

func safeLoginNext(request *http.Request) string {
	next := strings.TrimSpace(request.URL.Query().Get("next"))
	if next == "" || len(next) > 2048 || !strings.HasPrefix(next, "/") ||
		strings.HasPrefix(next, "//") || strings.ContainsAny(next, "\\\r\n") {
		return "/"
	}
	parsed, err := url.ParseRequestURI(next)
	if err != nil || parsed.IsAbs() || parsed.Host != "" || parsed.Path == "/login" || strings.HasPrefix(parsed.Path, "/login/") {
		return "/"
	}
	return next
}

func preserveTenantCookies(request *http.Request) []*http.Cookie {
	var preserved []*http.Cookie
	for _, cookie := range request.Cookies() {
		if cookie.Name == tenantSessionCookie || cookie.Name == tenantProfileCookie {
			preserved = append(preserved, cookie)
		}
	}
	return preserved
}

func privateWorkspaceReadAllowed(method, path string) bool {
	if method != http.MethodGet {
		return false
	}
	switch path {
	case "/api/workspaces", "/api/workspaces/suggest", "/api/file", "/api/file/raw", "/api/folder/download",
		"/api/logs", "/api/rollback/list", "/api/rollback/diff":
		return true
	default:
		return false
	}
}

func pathDenied(method, path string) bool {
	// Each Keycloak identity is already pinned to a dedicated WebUI and PVC.
	// Let that user discover and read files inside the backend-validated private
	// Home workspace, while keeping workspace registration/mutation and escape
	// APIs behind the chat administration boundary.
	if privateWorkspaceReadAllowed(method, path) {
		return false
	}
	for _, prefix := range deniedPrefixes {
		if path == prefix || strings.HasPrefix(path, prefix+"/") {
			return true
		}
	}
	return false
}

func (router *tenantRouter) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if request.URL.Path == "/healthz" {
		writer.Header().Set("Content-Type", "text/plain")
		_, _ = writer.Write([]byte("ok\n"))
		return
	}
	if request.Method == http.MethodGet && request.URL.Path == "/sw.js" {
		writer.Header().Set("Content-Type", "application/javascript; charset=utf-8")
		writer.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate")
		writer.Header().Set("Clear-Site-Data", `"cache", "storage"`)
		writer.Header().Set("Service-Worker-Allowed", "/")
		_, _ = writer.Write([]byte(retiredServiceWorker))
		return
	}
	subject := authenticatedSubject(request)
	if subject == "" {
		http.Error(writer, "authenticated identity required", http.StatusUnauthorized)
		return
	}
	// Keycloak is the public login authority. Never expose the tenant WebUI's
	// redundant password screen after a successful SSO handoff or when a user
	// follows a stale /login bookmark.
	if request.Method == http.MethodGet && request.URL.Path == "/login" {
		http.Redirect(writer, request, safeLoginNext(request), http.StatusFound)
		return
	}
	if router.serveTelegramWeb(writer, request, subject) {
		return
	}
	if pathDenied(request.Method, request.URL.Path) {
		http.Error(writer, "chat administration is disabled", http.StatusForbidden)
		return
	}
	slot, err := router.slotFor(subject)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusServiceUnavailable)
		return
	}
	target, err := url.Parse(router.backendURL(slot))
	if err != nil {
		http.Error(writer, "invalid tenant backend", http.StatusInternalServerError)
		return
	}
	proxy := httputil.NewSingleHostReverseProxy(target)
	proxy.ModifyResponse = injectChatBridge
	proxy.ErrorHandler = func(writer http.ResponseWriter, _ *http.Request, proxyErr error) {
		log.Printf("tenant slot %d unavailable", slot)
		http.Error(writer, "private chat runtime is starting", http.StatusBadGateway)
	}
	originalDirector := proxy.Director
	proxy.Director = func(outbound *http.Request) {
		originalDirector(outbound)
		cookies := preserveTenantCookies(outbound)
		for _, header := range []string{
			"Authorization", "Cookie", "X-Auth-Request-Access-Token",
			"X-Auth-Request-Email", "X-Auth-Request-Groups", "X-Auth-Request-User",
			"X-Forwarded-Access-Token", "X-Forwarded-Email", "X-Forwarded-Groups",
			"X-Forwarded-Preferred-Username", "X-Forwarded-User",
		} {
			outbound.Header.Del(header)
		}
		// Each Keycloak subject is already mapped to exactly one isolated pod.
		// Give that pod a non-sensitive, router-asserted identity while retaining
		// only its own WebUI cookies; OAuth credentials never reach the backend.
		outbound.Header.Set(trustedTenantHeader, fmt.Sprintf("slot-%d", slot))
		for _, cookie := range cookies {
			outbound.AddCookie(cookie)
		}
		outbound.Header.Del("Accept-Encoding")
	}
	proxy.ServeHTTP(writer, request)
}

func main() {
	slots, err := strconv.Atoi(os.Getenv("TENANT_SLOTS"))
	if err != nil || slots < 1 {
		log.Fatal("TENANT_SLOTS must be a positive integer")
	}
	statePath := os.Getenv("TENANT_STATE_PATH")
	if statePath == "" {
		statePath = "/state/tenants.json"
	}
	router, err := newTenantRouter(statePath, slots, func(slot int) string {
		return fmt.Sprintf("http://hermes-chat-tenant-%d.hermes-chat-tenant.hermes.svc.cluster.local:8787", slot)
	})
	if err != nil {
		log.Fatal(err)
	}
	router.backendAPIURL = func(slot int) string {
		return fmt.Sprintf("http://hermes-chat-tenant-%d.hermes-chat-tenant.hermes.svc.cluster.local:8642", slot)
	}
	telegramConfig, err := readTelegramConfig(os.Getenv("TELEGRAM_CONFIG_PATH"))
	if err != nil {
		log.Printf("Telegram is disabled: configuration is unavailable")
	} else if telegramConfig.BotToken != "" && telegramConfig.RelayKey != "" {
		router.telegram = newTelegramBot(telegramConfig, router)
		go router.telegram.run()
		log.Printf("Telegram transport enabled")
	} else {
		log.Printf("Telegram is prepared but disabled until bot_token is set in Vault")
	}
	server := &http.Server{
		Addr:              ":8080",
		Handler:           router,
		ReadHeaderTimeout: 10 * time.Second,
	}
	log.Printf("Hermes chat tenant router ready with %d isolated slots", slots)
	log.Fatal(server.ListenAndServe())
}
