// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// openai-proxy exposes an OpenAI-compatible HTTP server backed by the CMaaS SDK.
package main

import (
	"context"
	"flag"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/dashscope/dashscope-confidential-maas/sdk"
)

func main() {
	endpoint := flag.String("endpoint", "", "CMaaS WebSocket endpoint (required)")
	model := flag.String("model", "", "Model name (required)")
	listen := flag.String("listen", ":8080", "HTTP listen address")
	workers := flag.Int("workers", 8, "Number of concurrent CMaaS client sessions (= max concurrency)")
	cacheTTL := flag.Duration("cache-ttl", 12*time.Hour, "Public key cache TTL (set 0 to disable)")
	flag.Parse()

	if *endpoint == "" {
		log.Fatal("endpoint flag is required")
	}
	if *model == "" {
		log.Fatal("model flag is required")
	}

	apiKey := os.Getenv("DASHSCOPE_API_KEY")
	if apiKey == "" {
		log.Fatal("DASHSCOPE_API_KEY environment variable is required")
	}

	ttl := *cacheTTL
	if ttl == 0 {
		ttl = -1 // disable cache
	}

	opts := &sdk.Options{
		Endpoint:       *endpoint,
		APIKey:         apiKey,
		Model:          *model,
		ReuseSession:   false,
		PubKeyCacheTTL: ttl,
	}

	log.Printf("Creating pool of %d CMaaS client sessions...", *workers)
	p, err := newPool(*workers, opts)
	if err != nil {
		log.Fatalf("Failed to create client pool: %v", err)
	}
	defer p.close()

	h := &handler{pool: p, model: *model}

	srv := &http.Server{
		Addr:         *listen,
		Handler:      h,
		ReadTimeout:  5 * time.Minute,
		WriteTimeout: 10 * time.Minute,
		IdleTimeout:  2 * time.Minute,
	}

	log.Printf("OpenAI-compatible proxy listening on %s (workers=%d)", *listen, *workers)
	log.Printf("CMaaS backend endpoint=%s model=%q reuse_session=false", *endpoint, *model)
	log.Printf("  POST /v1/chat/completions (for --backend openai-chat)")
	log.Printf("  POST /v1/completions      (for --backend vllm / openai)")
	log.Printf("  GET  /v1/models")

	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("Server error: %v", err)
	}

	// Graceful shutdown
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	srv.Shutdown(ctx)
}
