package telnetproxy

import (
    "bufio"
    "context"
    "errors"
    "fmt"
    "io"
    "log"
    "net"
    "net/http"
    "strings"
    "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
}

// ServeHTTP implements caddyhttp.MiddlewareHandler.
func (tp TelnetProxy) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
    upgrader := websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            // Allow all origins (you might want to restrict this in production)
            return true
        },
    }

    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Printf("WebSocket upgrade error: %v", err)
        return err
    }
    defer ws.Close()

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

    return nil // Don't let caddy return a 404 after us
}

func (tp TelnetProxy) handleWebSocket(ws *websocket.Conn) {
    defer func() {
        if r := recover(); r != nil {
            log.Printf("Recovered from panic: %v", r)
        }
    }()

    var (
        hostPort string
        clientSecret string
        telnetConn net.Conn
        err error
    )

    // Initial Authentication and Host/Port
    err = ws.ReadJSON(map[string]string{
        "secret":     "",
        "host_port":  "",
    })
    if err != nil {
        log.Printf("Error reading initial JSON: %v", err)
        ws.WriteMessage(websocket.TextMessage, []byte("Error reading initial JSON"))
        return
    }


    err = ws.ReadJSON(&map[string]string{
        "secret": &clientSecret,
        "host_port": &hostPort,
    })
    if err != nil {
        log.Printf("Error reading initial JSON: %v", err)
        ws.WriteMessage(websocket.TextMessage, []byte("Error reading initial JSON"))
        return
    }



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


    // Telnet connection
    telnetConn, err = net.Dial("tcp", hostPort)
    if err != nil {
        log.Printf("Error connecting to Telnet server: %v", err)
        ws.WriteMessage(websocket.TextMessage, []byte("Error connecting to Telnet server: "+err.Error()))
        return
    }
    defer telnetConn.Close()

    // Channel for errors
    errChan := make(chan error, 2) // Buffered channel to avoid blocking

    // Copy from Telnet -> WebSocket (Read from telnet, send to client)
    go func() {
        // Using a buffered reader in case the server sends large chunks of text
        reader := bufio.NewReader(telnetConn)
        for {
            data, err := reader.ReadBytes('\n') // Read until a newline
             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 // End goroutine on EOF or error
            }

            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
            }

        }
    }()

    // Copy from WebSocket -> Telnet (Read from client, send to server)
    go func() {
        for {
            _, msg, err := ws.ReadMessage()
            if err != nil {
                 if !websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure){
                    log.Printf("Websocket read error (from browser): %v", err)
                     errChan <- fmt.Errorf("Websocket read error (from browser): %v", err)
                 }
                return
            }

            _, err = telnetConn.Write(msg)
            if err != nil {
                log.Printf("Telnet write error: %v", err)
                errChan <- fmt.Errorf("Telnet write error: %v", err)
                 return
            }
        }
    }()

    // Keep connection alive by waiting for an error on either read or write channel
    select {
    case <-errChan: // First Error closes connections
        log.Println("Closing connection")
        telnetConn.Close()
        ws.Close()

    case <-time.After(120 * time.Second): // Optional timeout.
        log.Println("Closing connection due to inactivity")
        telnetConn.Close()
        ws.Close()
    }
}