package mpd

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"
)

// KIDResolver is implemented by rewriters that can recover the live channel from
// a KID alone. The local server uses it so a rewrite still works after the
// upstream (ad-stitched) manifest session has expired and returns 410.
type KIDResolver interface {
	ResolveKID(kidHex string) (channel, originBase string, ok bool)
}

// Primer is implemented by rewriters that want to be told the fallback channel
// and origin discovered on earlier requests.
type Primer interface {
	Prime(channel, originBase string)
}

// LocalServer serves a rewritten live MPD on 127.0.0.1 for a downloader to
// refresh. Every GET re-fetches upstream and re-runs the rewriter.
type LocalServer struct {
	URL      string
	Upstream string
	Key      string // KID:KEY
	Headers  map[string]string

	rw         Rewriter
	kidHint    string
	ln         net.Listener
	httpServer *http.Server
}

// StartLocal starts the rewrite proxy. rw must not be nil; pass Passthrough{}
// to serve upstream unchanged.
func StartLocal(upstream, key string, headers map[string]string, rw Rewriter) (*LocalServer, error) {
	if rw == nil {
		rw = Passthrough{}
	}
	s := &LocalServer{Upstream: upstream, Key: key, Headers: headers, rw: rw}
	if key != "" && strings.Contains(key, ":") {
		s.kidHint = strings.SplitN(key, ":", 2)[0]
		if r, ok := rw.(KIDResolver); ok {
			if ch, base, found := r.ResolveKID(KIDHex(s.kidHint)); found {
				if p, ok := rw.(Primer); ok {
					p.Prime(ch, base)
				}
			}
		}
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	s.ln = ln
	mux := http.NewServeMux()
	mux.HandleFunc("/", s.handle)
	s.httpServer = &http.Server{Handler: mux}
	s.URL = fmt.Sprintf("http://%s/manifest.mpd", ln.Addr().String())
	go func() { _ = s.httpServer.Serve(ln) }()
	return s, nil
}

func (s *LocalServer) Close() {
	if s.httpServer != nil {
		_ = s.httpServer.Close()
	}
	if s.ln != nil {
		_ = s.ln.Close()
	}
}

func (s *LocalServer) handle(w http.ResponseWriter, r *http.Request) {
	body := ""
	req, err := http.NewRequest(http.MethodGet, s.Upstream, nil)
	if err == nil {
		for k, v := range s.Headers {
			req.Header.Set(k, v)
		}
		client := &http.Client{Timeout: 20 * time.Second}
		if resp, err := client.Do(req); err == nil {
			b, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
			_ = resp.Body.Close()
			if resp.StatusCode >= 200 && resp.StatusCode < 300 && len(b) > 0 {
				body = string(b)
			}
		}
	}
	// Upstream being gone is not fatal: a rewriter with a KID fallback can still
	// publish the encoder timeline so the downloader keeps pulling segments.
	payload, err := s.rw.Rewrite([]byte(body), s.kidHint)
	if err != nil {
		http.Error(w, err.Error(), 502)
		return
	}
	w.Header().Set("Content-Type", "application/dash+xml")
	w.WriteHeader(200)
	_, _ = w.Write(payload)
}