package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"
)

func main() {
	version := "1.0.2"
	targetUrl := "https://api.openai.com/" // 目标域名和端口
	target, err := url.Parse(targetUrl)
	if err != nil {
		log.Fatal(err)
	}

	// 创建反向代理
	proxy := httputil.NewSingleHostReverseProxy(target)
	// 设置输出缓冲区的刷新时间间隔
	proxy.FlushInterval = 1 * time.Second

	// 修改请求头，将Host设置为目标域名
	proxy.Director = func(req *http.Request) {
		req.Host = target.Host
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
	}

	// 设置超时时间
	timeout := time.Duration(1 * time.Minute)
	proxy.Transport = &http.Transport{
		Proxy:                 http.ProxyFromEnvironment,
		DialContext:           (&net.Dialer{Timeout: timeout}).DialContext,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}

	// 打印HTTP请求和响应的日志
	proxy.ModifyResponse = func(resp *http.Response) error {
		// 打印HTTP响应的日志
		responseDump, err := httputil.DumpResponse(resp, true)
		if err != nil {
			log.Printf("Failed to dump response: %v\n", err)
		} else {
			log.Printf("%s Response: \n%s\n", time.Now().Format("2006-01-02 15:04:05"), string(responseDump))
		}

		return nil
	}

	// 设置日志前缀和输出位置
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	// 启动HTTP服务器
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// 设置 Content-Type 和 Transfer-Encoding
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Transfer-Encoding", "chunked")

		// 将http.ResponseWriter转换为http.Flusher接口
		flusher, ok := w.(http.Flusher)
		if !ok {
			http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
			log.Printf("Streaming unsupported!")
			return
		} else {
			log.Printf("Streaming supported!")
		}

		// 打印HTTP请求日志
		requestDump, err := httputil.DumpRequest(r, true)
		if err != nil {
			log.Printf("Failed to dump request: \n%v\n", err)
		} else {
			log.Printf("proxy version: %s\n%s Request: %s\n", version, time.Now().Format("2006-01-02 15:04:05"), string(requestDump))
		}
		// 反向代理转发
		proxy.ServeHTTP(w, r)

		// 在每个数据块发送完毕后刷新响应缓冲区
		flusher.Flush()
	})

	http.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
		// 设置响应头
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		// 设置Transfer-Encoding字段为"chunked"
		w.Header().Set("Transfer-Encoding", "chunked")
		w.WriteHeader(http.StatusOK) // 写入响应头

		// 发送初始化事件
		fmt.Fprint(w, "event: connect\n")
		fmt.Fprint(w, ": server connected\n\n")
		w.(http.Flusher).Flush()

		// 发送数据
		for i := 0; i < 10; i++ {
			fmt.Fprintf(w, "data: {\"count\": %d}\n\n", i)
			w.(http.Flusher).Flush()
			time.Sleep(time.Second)
		}

		// 结束 SSE 连接
		fmt.Fprint(w, "event: close\n")
		w.(http.Flusher).Flush()
		time.Sleep(time.Second)
		return
	})

	log.Printf("Starting server on port 9000...\n")
	if err := http.ListenAndServe(":9000", nil); err != nil {
		log.Fatal(err)
	}
}