package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

type Request struct {
	ID      uint64              `json:"id"`
	Method  string              `json:"method"`
	Path    string              `json:"path"`
	Headers map[string][]string `json:"headers"`
	Body    []byte              `json:"body"`
}

type Response struct {
	ID      uint64              `json:"id"`
	Status  int                 `json:"status"`
	Headers map[string][]string `json:"headers"`
	Body    []byte              `json:"body"`
}

type Tunnel struct {
	conn     *websocket.Conn
	mu       sync.Mutex
	pending  map[uint64]chan Response
	pendMu   sync.RWMutex
	counter  uint64
	lastPing time.Time
}

var (
	tunnel   *Tunnel
	tunnelMu sync.RWMutex
	upgrader = websocket.Upgrader{
		CheckOrigin:     func(r *http.Request) bool { return true },
		ReadBufferSize:  65536,
		WriteBufferSize: 65536,
	}
	authToken = "your-secret-token-change-me"
)

func main() {
	http.HandleFunc("/", handler)
	log.Println("Server starting on :7860")
	log.Fatal(http.ListenAndServe(":7860", nil))
}

func handler(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/_tunnel":
		handleTunnel(w, r)
	case "/_health":
		handleHealth(w)
	default:
		handleProxy(w, r)
	}
}

func handleHealth(w http.ResponseWriter) {
	tunnelMu.RLock()
	connected := tunnel != nil
	tunnelMu.RUnlock()

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]any{
		"status":    "ok",
		"connected": connected,
		"time":      time.Now().Unix(),
	})
}

func handleTunnel(w http.ResponseWriter, r *http.Request) {
	if r.Header.Get("X-Tunnel-Token") != authToken {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("upgrade failed: %v", err)
		return
	}

	t := &Tunnel{
		conn:     conn,
		pending:  make(map[uint64]chan Response),
		lastPing: time.Now(),
	}

	tunnelMu.Lock()
	if tunnel != nil {
		tunnel.conn.Close()
	}
	tunnel = t
	tunnelMu.Unlock()

	log.Println("client connected")
	t.readLoop()

	tunnelMu.Lock()
	if tunnel == t {
		tunnel = nil
	}
	tunnelMu.Unlock()
	log.Println("client disconnected")
}

func (t *Tunnel) readLoop() {
	defer t.conn.Close()

	for {
		_, data, err := t.conn.ReadMessage()
		if err != nil {
			return
		}

		var resp Response
		if json.Unmarshal(data, &resp) != nil {
			continue
		}

		t.pendMu.RLock()
		ch, ok := t.pending[resp.ID]
		t.pendMu.RUnlock()

		if ok {
			select {
			case ch <- resp:
			default:
			}
		}
	}
}

func (t *Tunnel) send(req Request) (Response, error) {
	ch := make(chan Response, 1)

	t.pendMu.Lock()
	t.pending[req.ID] = ch
	t.pendMu.Unlock()

	defer func() {
		t.pendMu.Lock()
		delete(t.pending, req.ID)
		t.pendMu.Unlock()
	}()

	data, _ := json.Marshal(req)

	t.mu.Lock()
	err := t.conn.WriteMessage(websocket.TextMessage, data)
	t.mu.Unlock()

	if err != nil {
		return Response{}, err
	}

	select {
	case resp := <-ch:
		return resp, nil
	case <-time.After(60 * time.Second):
		return Response{Status: 504}, nil
	}
}

func handleProxy(w http.ResponseWriter, r *http.Request) {
	tunnelMu.RLock()
	t := tunnel
	tunnelMu.RUnlock()

	if t == nil {
		http.Error(w, "tunnel not connected", http.StatusBadGateway)
		return
	}

	body, _ := io.ReadAll(r.Body)

	req := Request{
		ID:      atomic.AddUint64(&t.counter, 1),
		Method:  r.Method,
		Path:    r.URL.RequestURI(),
		Headers: r.Header,
		Body:    body,
	}

	resp, err := t.send(req)
	if err != nil {
		http.Error(w, "tunnel error", http.StatusBadGateway)
		return
	}

	for k, vals := range resp.Headers {
		for _, v := range vals {
			w.Header().Add(k, v)
		}
	}

	w.WriteHeader(resp.Status)
	w.Write(resp.Body)
}