package main

import (
	"context"
	"crypto/subtle"
	"embed"
	"encoding/json"
	"errors"
	"flag"
	"io/fs"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
)

//go:embed web/*
var assets embed.FS

type server struct {
	cfg           config
	token, origin string
	slots         chan struct{}
	newProvider   func(string) (asrProvider, error)
	sessionLimit  time.Duration
}

type modelOption struct {
	ID      string `json:"id"`
	Enabled bool   `json:"enabled"`
}

func (s *server) bailianCredentials(model string) (string, string) {
	if strings.HasPrefix(model, "fun-") {
		return s.cfg.funASRKey, s.cfg.funASRBase
	}
	return s.cfg.bailianKey, s.cfg.bailianBase
}

func (s *server) models() []modelOption {
	out := []modelOption{}
	for _, id := range strings.Split(env("BAILIAN_ASR_MODELS", "fun-asr-realtime"), ",") {
		id = strings.TrimSpace(id)
		if id != "" && id != "volc-bigmodel" {
			key, _ := s.bailianCredentials(id)
			out = append(out, modelOption{id, key != ""})
		}
	}
	return append(out, modelOption{"volc-bigmodel", s.cfg.volcAppKey != "" || s.cfg.volcAppID != "" && s.cfg.volcToken != ""})
}

func (s *server) provider(id string) (asrProvider, error) {
	for _, model := range s.models() {
		if model.ID != id {
			continue
		}
		if !model.Enabled {
			return nil, errors.New("provider credentials are not configured")
		}
		if id == "volc-bigmodel" {
			return newVolcASR(agentCfg{VolcAppKey: s.cfg.volcAppKey, VolcAppID: s.cfg.volcAppID, VolcAccessToken: s.cfg.volcToken, VolcResourceID: s.cfg.volcResourceID, VolcWS: s.cfg.volcWS}), nil
		}
		key, endpoint := s.bailianCredentials(id)
		return newBailianASR(id, agentCfg{BailianKey: key, BailianWssBaseURL: endpoint}), nil
	}
	return nil, errors.New("unknown ASR model")
}

func (s *server) authorized(r *http.Request) bool {
	token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
	for _, protocol := range websocket.Subprotocols(r) {
		if strings.HasPrefix(protocol, "auth.") {
			token = strings.TrimPrefix(protocol, "auth.")
		}
	}
	return len(token) == len(s.token) && subtle.ConstantTimeCompare([]byte(token), []byte(s.token)) == 1
}

func (s *server) sameOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		return true
	} // CLI clients still need the access token.
	if s.origin != "" {
		return origin == s.origin
	}
	u, err := url.Parse(origin)
	return err == nil && u.User == nil && u.Path == "" && u.RawQuery == "" && u.Fragment == "" && (u.Scheme == "http" || u.Scheme == "https") && u.Host == r.Host
}

func (s *server) handler(ctx context.Context) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"status":"ok","scope":"asr-only","llm":false,"tts":false}`))
	})
	mux.HandleFunc("GET /api/models", func(w http.ResponseWriter, r *http.Request) {
		if !s.authorized(r) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(s.models())
	})
	mux.HandleFunc("GET /ws", func(w http.ResponseWriter, r *http.Request) { s.serveWS(ctx, w, r) })
	web, _ := fs.Sub(assets, "web")
	mux.Handle("GET /", http.FileServerFS(web))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("Referrer-Policy", "no-referrer")
		w.Header().Set("Cache-Control", "no-store")
		w.Header().Set("Content-Security-Policy", "default-src 'self'; script-src 'self'; style-src 'self'; connect-src 'self'; frame-ancestors 'none'; base-uri 'none'; form-action 'self'")
		mux.ServeHTTP(w, r)
	})
}

func (s *server) serveWS(parent context.Context, w http.ResponseWriter, r *http.Request) {
	if !s.authorized(r) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	if !s.sameOrigin(r) {
		http.Error(w, "origin rejected", http.StatusForbidden)
		return
	}
	select {
	case s.slots <- struct{}{}:
		defer func() { <-s.slots }()
	default:
		http.Error(w, "session limit reached", http.StatusTooManyRequests)
		return
	}
	up := websocket.Upgrader{Subprotocols: []string{"asr.v1"}, CheckOrigin: s.sameOrigin, HandshakeTimeout: 5 * time.Second}
	conn, err := up.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	limit := s.sessionLimit
	if limit == 0 {
		limit = 5 * time.Minute
	}
	ctx, cancel := context.WithTimeout(parent, limit)
	defer cancel()
	defer conn.Close()
	conn.SetReadLimit(64 << 10)
	stopClose := context.AfterFunc(ctx, func() { _ = conn.Close() })
	defer stopClose()
	var writeMu sync.Mutex
	var stateMu sync.Mutex
	var active *asrRun
	send := func(expected *asrRun, value any) {
		writeMu.Lock()
		defer writeMu.Unlock()
		stateMu.Lock()
		valid := expected == nil || active == expected
		stateMu.Unlock()
		if !valid || ctx.Err() != nil {
			return
		}
		_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
		if err := conn.WriteJSON(value); err != nil {
			cancel()
		}
	}
	stop := func() {
		stateMu.Lock()
		old := active
		active = nil
		stateMu.Unlock()
		if old != nil {
			old.cancel()
			old.provider.close()
		}
	}
	defer stop()
	send(nil, map[string]any{"type": "ready", "sample_rate": 16000, "channels": 1, "format": "pcm_s16le", "max_session_seconds": int(limit.Seconds()), "llm": false, "tts": false})
	for {
		kind, data, err := conn.ReadMessage()
		if err != nil {
			return
		}
		if kind == websocket.BinaryMessage {
			stateMu.Lock()
			run := active
			stateMu.Unlock()
			if run == nil {
				send(nil, map[string]string{"type": "error", "error": "start ASR before sending audio"})
				continue
			}
			if len(data)%2 != 0 {
				send(run, map[string]string{"type": "error", "error": "PCM16 must have an even byte length"})
				continue
			}
			if err := run.provider.sendAudio(data); err != nil {
				send(run, map[string]string{"type": "error", "error": "upstream audio write failed"})
				stop()
			}
			continue
		}
		var msg struct {
			Type  string `json:"type"`
			Model string `json:"model"`
		}
		if kind != websocket.TextMessage || json.Unmarshal(data, &msg) != nil {
			send(nil, map[string]string{"type": "error", "error": "invalid control message"})
			continue
		}
		switch msg.Type {
		case "start":
			stop()
			factory := s.newProvider
			if factory == nil {
				factory = s.provider
			}
			p, err := factory(msg.Model)
			if err != nil {
				send(nil, map[string]string{"type": "error", "error": err.Error()})
				continue
			}
			runCtx, runCancel := context.WithCancel(ctx)
			run := &asrRun{provider: p, cancel: runCancel, ctx: runCtx}
			if err := p.start(runCtx); err != nil {
				runCancel()
				p.close()
				send(nil, map[string]string{"type": "error", "error": "ASR start failed; check provider credentials, model and endpoint"})
				continue
			}
			stateMu.Lock()
			active = run
			stateMu.Unlock()
			send(run, map[string]string{"type": "asr-started", "model": msg.Model})
			go func() {
				defer runCancel()
				defer p.close()
				for {
					select {
					case <-runCtx.Done():
						return
					case ev, ok := <-p.events():
						if !ok {
							send(run, map[string]string{"type": "asr-stopped"})
							stateMu.Lock()
							if active == run {
								active = nil
							}
							stateMu.Unlock()
							return
						}
						send(run, map[string]string{"type": ev.Typ, "text": ev.Text, "code": ev.Code, "error": ev.Error})
					}
				}
			}()
		case "finish":
			stateMu.Lock()
			run := active
			stateMu.Unlock()
			if run != nil {
				run.finishOnce.Do(func() {
					if err := run.provider.finish(); err != nil {
						send(run, map[string]string{"type": "error", "error": "upstream finish failed"})
						stop()
						return
					}
					go func() {
						select {
						case <-run.ctx.Done():
						case <-time.After(10 * time.Second):
							stateMu.Lock()
							same := active == run
							stateMu.Unlock()
							if same {
								run.cancel()
								run.provider.close()
								send(run, map[string]string{"type": "asr-stopped"})
							}
						}
					}()
				})
			}
		case "cancel":
			stop()
			send(nil, map[string]string{"type": "asr-stopped"})
		default:
			send(nil, map[string]string{"type": "error", "error": "unsupported command; LLM/TTS are disabled"})
		}
	}
}

type asrRun struct {
	provider   asrProvider
	cancel     context.CancelFunc
	ctx        context.Context
	finishOnce sync.Once
}

func main() {
	health := flag.Bool("healthcheck", false, "check local process liveness")
	flag.Parse()
	if *health {
		client := http.Client{Timeout: 2 * time.Second}
		resp, err := client.Get("http://127.0.0.1:8080/healthz")
		if err != nil {
			os.Exit(1)
		}
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			os.Exit(1)
		}
		return
	}
	cfg := loadConfig()
	token := os.Getenv("ASR_WEB_TOKEN")
	origin := os.Getenv("PUBLIC_ORIGIN")
	if err := validateConfig(token, origin, cfg); err != nil {
		log.Fatal(err)
	}
	s := &server{cfg: cfg, token: token, origin: origin, slots: make(chan struct{}, 4)}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	srv := &http.Server{Addr: env("ASR_LISTEN", "127.0.0.1:8080"), Handler: s.handler(ctx), ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 60 * time.Second, MaxHeaderBytes: 16 << 10}
	go func() {
		<-ctx.Done()
		deadline, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		_ = srv.Shutdown(deadline)
	}()
	log.Printf("ASR-only service listening on %s; LLM/TTS disabled", srv.Addr)
	if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
		log.Fatal(err)
	}
}