package telnetproxy

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/gorilla/websocket"
)

// TelnetProxy is a Caddy module that proxies WebSocket connections to Telnet servers.
type TelnetProxy struct {
	Secret string `json:"secret,omitempty"`
}

// CaddyModule returns the Caddy module information.
func (TelnetProxy) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.telnet_proxy",
		New: func() caddy.Module { return new(TelnetProxy) },
	}
}

// UnmarshalCaddyfile implements caddyfile.Unmarshaler.
func (tp *TelnetProxy) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		if !d.Args(&tp.Secret) {
			return d.ArgErr()
		}
	}
	return nil
}

// authMessage is used for the initial authentication and connection details.
type authMessage struct {
	Secret   string `json:"secret"`
	HostPort string `json:"host_port"`
}

// ServeHTTP implements caddyhttp.MiddlewareHandler.
func (tp TelnetProxy) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
	upgrader := websocket.Upgrader{
		// In production, restrict allowed origins to trusted hosts.
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return err
	}
	// The connection will be closed in handleWebSocket.
	defer ws.Close()

	// Handle the WebSocket connection in a separate goroutine.
	go tp.handleWebSocket(ws)

	// Returning nil prevents further Caddy middleware from interfering.
	return nil
}

func (tp TelnetProxy) handleWebSocket(ws *websocket.Conn) {
	// Recover from panics to avoid crashing the server.
	defer func() {
		if r := recover(); r != nil {
			log.Printf("Recovered from panic: %v", r)
		}
	}()

	// Read and decode the initial authentication message.
	var msg authMessage
	if err := ws.ReadJSON(&msg); err != nil {
		log.Printf("Error reading initial JSON: %v", err)
		ws.WriteMessage(websocket.TextMessage, []byte("Error reading initial JSON"))
		return
	}

	// Check authentication.
	if msg.Secret != tp.Secret {
		log.Printf("Authentication failed. Received secret: %s", msg.Secret)
		ws.WriteMessage(websocket.TextMessage, []byte("Authentication failed"))
		return
	}

	// Connect to the Telnet server with a timeout.
	telnetConn, err := net.DialTimeout("tcp", msg.HostPort, 10*time.Second)
	if err != nil {
		log.Printf("Error connecting to Telnet server at %s: %v", msg.HostPort, err)
		ws.WriteMessage(websocket.TextMessage, []byte("Error connecting to Telnet server: "+err.Error()))
		return
	}
	defer telnetConn.Close()

	// Create a channel to signal errors from either direction.
	errChan := make(chan error, 2)

	// Goroutine: Telnet -> WebSocket.
	go func() {
		reader := bufio.NewReader(telnetConn)
		for {
			// Read until newline; adjust if the Telnet protocol doesn't send newline-terminated data.
			data, err := reader.ReadBytes('\n')
			if err != nil {
				if !errors.Is(err, io.EOF) {
					log.Printf("Telnet read error: %v", err)
					errChan <- fmt.Errorf("Telnet read error: %v", err)
				}
				return
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, data); err != nil {
				log.Printf("WebSocket write error (from Telnet): %v", err)
				errChan <- fmt.Errorf("WebSocket write error (from Telnet): %v", err)
				return
			}
		}
	}()

	// Goroutine: WebSocket -> Telnet.
	go func() {
		for {
			_, data, err := ws.ReadMessage()
			if err != nil {
				if !websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
					log.Printf("WebSocket read error (from client): %v", err)
					errChan <- fmt.Errorf("WebSocket read error (from client): %v", err)
				}
				return
			}
			if _, err := telnetConn.Write(data); err != nil {
				log.Printf("Telnet write error: %v", err)
				errChan <- fmt.Errorf("Telnet write error: %v", err)
				return
			}
		}
	}()

	// Wait until an error occurs or a timeout happens.
	select {
	case <-errChan:
		log.Println("Closing connection due to an error.")
	case <-time.After(120 * time.Second):
		log.Println("Closing connection due to inactivity.")
	}

	// Close both connections (deferred calls will handle this if not already closed).
	telnetConn.Close()
	ws.Close()
}