package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"
)

// 配置常量
const (
	BasePath   = "/edge-aac/903hd"
	ServerPort = ":8080"
)

// ProxyServer 代理服务器结构 (帶狀態管理)
type ProxyServer struct {
	jar              *cookiejar.Jar
	cookieExpiration time.Time
	upstreamHost     string
	mu               sync.RWMutex
}

// NewProxyServer 创建新的代理服务器
func NewProxyServer() *ProxyServer {
	jar, _ := cookiejar.New(nil)
	return &ProxyServer{
		jar:              jar,
		cookieExpiration: time.Now(), // 設置為當前時間，以觸發首次請求時立即獲取
	}
}

// ensureValidCookies 確保 Cookie 有效，如果過期則刷新
func (ps *ProxyServer) ensureValidCookies() error {
	ps.mu.RLock()
	isExpired := time.Now().After(ps.cookieExpiration)
	ps.mu.RUnlock()

	if isExpired {
		log.Println("Cookie expired or not yet fetched, refreshing...")
		ps.mu.Lock()
		defer ps.mu.Unlock()

		if time.Now().After(ps.cookieExpiration) {
			if err := ps.refreshCookies(); err != nil {
				log.Printf("Failed to refresh cookies: %v", err)
				return err
			}
			log.Println("Cookies refreshed successfully.")
		}
	}
	return nil
}

// refreshCookies 執行完整的 Cookie 獲取和解析流程 (已按新邏輯修改)
func (ps *ProxyServer) refreshCookies() error {
	// --- 步驟 1 & 2: 獲取 liveJsUrl ---
	liveJsURL, err := getLiveJsURL("https://www.881903.com/live/903")
	if err != nil {
		return err
	}

	// --- 步驟 3: 獲取 302 回應 ---
	resp302, err := get302Response(liveJsURL)
	if err != nil {
		return err
	}
	defer resp302.Body.Close()

	// --- 步驟 4: 訪問重定向地址以獲取真正的授權 Cookie ---
	redirectURL, err := resp302.Location()
	if err != nil {
		return fmt.Errorf("failed to get redirect location: %w", err)
	}

	respFromRedirect, err := http.Get(redirectURL.String())
	if err != nil {
		return fmt.Errorf("failed to visit redirect URL: %w", err)
	}
	defer respFromRedirect.Body.Close()

	cookies := respFromRedirect.Cookies()
	if len(cookies) == 0 {
		return fmt.Errorf("no cookies found in redirect response")
	}

	// --- 按照新邏輯設置狀態 ---

	// 1. 清空舊的 Jar 並設置新的 Cookies
	ps.jar, _ = cookiejar.New(nil)
	ps.jar.SetCookies(redirectURL, cookies)

	// 2. 將 Upstream Host 設置為 redirectURL 的 Host
	ps.upstreamHost = redirectURL.Host
	log.Printf("Upstream host set to: %s", ps.upstreamHost)

	// 3. 將過期時間設置為當前時間 + 24 小時
	ps.cookieExpiration = time.Now().Add(24 * time.Hour)
	log.Printf("New cookie expiration set to: %s", ps.cookieExpiration.Format(time.RFC3339))

	return nil
}

// proxyRequest 處理代理請求的核心邏輯
func (ps *ProxyServer) proxyRequest(w http.ResponseWriter, r *http.Request) {
	if err := ps.ensureValidCookies(); err != nil {
		http.Error(w, fmt.Sprintf("Internal Server Error: %v", err), http.StatusInternalServerError)
		return
	}

	ps.mu.RLock()
	upstreamHost := ps.upstreamHost
	ps.mu.RUnlock()

	targetURL, err := ps.buildTargetURL(r.URL.Path, upstreamHost)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	log.Printf("Proxying request for %s to %s", r.URL.Path, targetURL)

	client := &http.Client{
		Jar:     ps.jar,
		Timeout: 30 * time.Second,
	}
	proxyReq, err := http.NewRequest(r.Method, targetURL, nil)
	if err != nil {
		http.Error(w, "Failed to create proxy request", http.StatusInternalServerError)
		return
	}

	resp, err := client.Do(proxyReq)
	if err != nil {
		http.Error(w, fmt.Sprintf("Proxy request failed: %v", err), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	enableCORS(w)
	for key, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		http.Error(w, "Failed to read response", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(resp.StatusCode)
	w.Write(body)
}

// --- ServeHTTP 和其他輔助函數 (與之前版本相同) ---
func (ps *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodOptions {
		enableCORS(w)
		w.WriteHeader(http.StatusNoContent)
		return
	}
	ps.proxyRequest(w, r)
}

func (ps *ProxyServer) buildTargetURL(path string, host string) (string, error) {
	switch {
	case path == "/" || path == "/playlist.m3u8":
		return fmt.Sprintf("https://%s%s/chunks.m3u8", host, BasePath), nil
	case strings.HasPrefix(path, "/chunk/"):
		return fmt.Sprintf("https://%s%s/%s", host, BasePath, strings.TrimPrefix(path, "/chunk/")), nil
	default:
		// 增加對 .aac 文件的直接代理
		if strings.HasSuffix(path, ".aac") {
			return fmt.Sprintf("https://%s%s%s", host, BasePath, path), nil
		}
		return "", fmt.Errorf("unsupported path: %s", path)
	}
}

func enableCORS(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
}

// --- Cookie 獲取輔助函數 (與之前版本相同) ---
func getLiveJsURL(pageURL string) (string, error) {
	log.Println("Fetching liveJsUrl from", pageURL)
	resp, err := http.Get(pageURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	re := regexp.MustCompile(`"liveJsUrl":"([^"]+)"`)
	matches := re.FindStringSubmatch(string(body))
	if len(matches) < 2 {
		return "", fmt.Errorf("liveJsUrl not found")
	}
	log.Println("Found liveJsUrl")
	return matches[1], nil
}

func get302Response(jsURL string) (*http.Response, error) {
	log.Println("Fetching 302 redirect from", jsURL)
	client := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	req, err := http.NewRequest("GET", jsURL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("accept", "*/*")
	req.Header.Set("referer", "https://www.881903.com/live/903")
	req.Header.Set("user-agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36")
	resp, err := client.Do(req)
	if urlErr, ok := err.(*url.Error); ok && urlErr.Err == http.ErrUseLastResponse {
		return resp, nil
	}
	return resp, err
}

// --- Main 函數 (與之前版本相同) ---
func main() {
	proxy := NewProxyServer()

	log.Println("Performing initial cookie fetch...")
	if err := proxy.ensureValidCookies(); err != nil {
		log.Fatalf("Initial cookie fetch failed: %v. Server cannot start.", err)
	}

	http.Handle("/", proxy)
	log.Printf("Starting HLS proxy server on port %s", ServerPort)
	log.Printf("Access playlist at: http://localhost%s/playlist.m3u8", ServerPort)
	if err := http.ListenAndServe(ServerPort, nil); err != nil {
		log.Fatal("Server failed to start:", err)
	}
}