package main

import (
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "net/url"
    "os"
    "strings"
    "time"

    "github.com/gorilla/websocket"
)

type WebSocketLoggingProxy struct {
    targetURL *url.URL
    upgrader  websocket.Upgrader
}

type ConnectionLog struct {
    Timestamp  time.Time `json:"timestamp"`
    Event      string    `json:"event"` // "connect", "disconnect"
    RemoteAddr string    `json:"remote_addr"`
    URL        string    `json:"url"`
    Headers    map[string][]string `json:"headers"`
}

type MessageLog struct {
    Timestamp  time.Time `json:"timestamp"`
    Direction  string    `json:"direction"` // "incoming" (от клиента), "outgoing" (к клиенту)
    Message    string    `json:"message"`
    MessageType int      `json:"message_type"` // 1=text, 2=binary
    RemoteAddr string    `json:"remote_addr"`
}

type LogEntry struct {
    Connection ConnectionLog `json:"connection"`
    Messages   []MessageLog  `json:"messages,omitempty"`
}

func NewWebSocketLoggingProxy(targetURL string) (*WebSocketLoggingProxy, error) {
    target, err := url.Parse(targetURL)
    if err != nil {
        return nil, fmt.Errorf("invalid target URL: %w", err)
    }
    
    // Проверяем, что схема правильная для WebSocket
    if target.Scheme != "ws" && target.Scheme != "wss" {
        return nil, fmt.Errorf("target URL must use ws:// or wss:// scheme")
    }
    
    upgrader := websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            // Разрешаем все источники для прокси
            return true
        },
    }
    
    return &WebSocketLoggingProxy{
        targetURL: target,
        upgrader:  upgrader,
    }, nil
}

func (wsp *WebSocketLoggingProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    // Логируем подключение
    connLog := ConnectionLog{
        Timestamp:  time.Now(),
        Event:      "connect",
        RemoteAddr: r.RemoteAddr,
        URL:        r.URL.String(),
        Headers:    r.Header,
    }
    wsp.printConnectionLog(connLog)
    
    // Обновляем HTTP соединение до WebSocket
    clientConn, err := wsp.upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Printf("Failed to upgrade connection: %v", err)
        return
    }
    defer clientConn.Close()
    
    // Формируем целевой URL для WebSocket
    targetWSURL := wsp.buildTargetURL(r)
    
    // Фильтруем заголовки, исключая те, которые конфликтуют с WebSocket протоколом
    targetHeaders := wsp.filterHeaders(r.Header)
    
    // Подключаемся к целевому WebSocket серверу
    targetConn, _, err := websocket.DefaultDialer.Dial(targetWSURL, targetHeaders)
    if err != nil {
        log.Printf("Failed to connect to target: %v", err)
        clientConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "Failed to connect to target"))
        return
    }
    defer targetConn.Close()
    
    // Создаем канал для сообщений
    messages := make([]MessageLog, 0)
    
    // Канал для завершения
    done := make(chan struct{})
    
    // Горутина для пересылки сообщений от клиента к целевому серверу
    go func() {
        defer close(done)
        for {
            messageType, message, err := clientConn.ReadMessage()
            if err != nil {
                if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                    log.Printf("Error reading from client: %v", err)
                }
                return
            }
            
            // Логируем входящее сообщение
            msgLog := MessageLog{
                Timestamp:   time.Now(),
                Direction:   "incoming",
                Message:     wsp.formatMessage(message, messageType),
                MessageType: messageType,
                RemoteAddr:  r.RemoteAddr,
            }
            messages = append(messages, msgLog)
            wsp.printMessageLog(msgLog)
            
            // Пересылаем сообщение целевому серверу
            if err := targetConn.WriteMessage(messageType, message); err != nil {
                log.Printf("Error writing to target: %v", err)
                return
            }
        }
    }()
    
    // Горутина для пересылки сообщений от целевого сервера к клиенту
    go func() {
        for {
            messageType, message, err := targetConn.ReadMessage()
            if err != nil {
                if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                    log.Printf("Error reading from target: %v", err)
                }
                return
            }
            
            // Логируем исходящее сообщение
            msgLog := MessageLog{
                Timestamp:   time.Now(),
                Direction:   "outgoing",
                Message:     wsp.formatMessage(message, messageType),
                MessageType: messageType,
                RemoteAddr:  r.RemoteAddr,
            }
            messages = append(messages, msgLog)
            wsp.printMessageLog(msgLog)
            
            // Пересылаем сообщение клиенту
            if err := clientConn.WriteMessage(messageType, message); err != nil {
                log.Printf("Error writing to client: %v", err)
                return
            }
        }
    }()
    
    // Ждем завершения одной из горутин
    <-done
    
    // Логируем отключение
    disconnectLog := ConnectionLog{
        Timestamp:  time.Now(),
        Event:      "disconnect",
        RemoteAddr: r.RemoteAddr,
        URL:        r.URL.String(),
        Headers:    r.Header,
    }
    wsp.printConnectionLog(disconnectLog)
    
    // Выводим итоговую запись с всеми сообщениями
    if len(messages) > 0 {
        logEntry := LogEntry{
            Connection: connLog,
            Messages:   messages,
        }
        wsp.printFullLog(logEntry)
    }
}

func (wsp *WebSocketLoggingProxy) buildTargetURL(r *http.Request) string {
    // Формируем целевой URL, сохраняя путь и query параметры
    targetPath := r.URL.Path
    if r.URL.RawQuery != "" {
        targetPath += "?" + r.URL.RawQuery
    }
    
    return fmt.Sprintf("%s://%s%s", wsp.targetURL.Scheme, wsp.targetURL.Host, targetPath)
}

func (wsp *WebSocketLoggingProxy) filterHeaders(originalHeaders http.Header) http.Header {
    // Создаем новый набор заголовков, исключая те, которые конфликтуют с WebSocket
    filtered := make(http.Header)
    
    // Заголовки, которые нужно исключить (они устанавливаются автоматически библиотекой)
    // Используем lowercase для case-insensitive сравнения
    excludedHeadersLower := map[string]bool{
        "upgrade":                 true,
        "connection":              true,
        "sec-websocket-key":       true,
        "sec-websocket-version":   true,
        "sec-websocket-extensions": true,
        "sec-websocket-protocol":   true,
    }
    
    // Копируем только разрешенные заголовки
    for k, v := range originalHeaders {
        // Преобразуем в lowercase для сравнения
        keyLower := strings.ToLower(k)
        if !excludedHeadersLower[keyLower] {
            // Сохраняем оригинальное имя заголовка
            filtered[k] = make([]string, len(v))
            copy(filtered[k], v)
        }
    }
    
    return filtered
}

func (wsp *WebSocketLoggingProxy) formatMessage(message []byte, messageType int) string {
    if messageType == websocket.TextMessage {
        // Пытаемся форматировать как JSON для лучшей читаемости
        // Сохраняем полное сообщение без обрезки
        var jsonData interface{}
        if err := json.Unmarshal(message, &jsonData); err == nil {
            // Форматируем JSON с отступами, сохраняя все данные полностью
            if formatted, err := json.MarshalIndent(jsonData, "", "  "); err == nil {
                return string(formatted)
            }
        }
        // Если не JSON или ошибка форматирования, возвращаем как есть (полностью)
        return string(message)
    } else {
        // Для бинарных сообщений показываем размер
        return fmt.Sprintf("[Binary: %d bytes]", len(message))
    }
}

func (wsp *WebSocketLoggingProxy) printConnectionLog(connLog ConnectionLog) {
    fmt.Println(strings.Repeat("=", 80))
    fmt.Printf("[%s] WebSocket %s\n", connLog.Timestamp.Format("2006-01-02 15:04:05"), strings.ToUpper(connLog.Event))
    fmt.Printf("Remote: %s\n", connLog.RemoteAddr)
    fmt.Printf("URL: %s\n", connLog.URL)
    if len(connLog.Headers) > 0 {
        fmt.Println("\nHeaders:")
        for k, v := range connLog.Headers {
            fmt.Printf("  %s: %s\n", k, strings.Join(v, ", "))
        }
    }
    fmt.Println(strings.Repeat("=", 80))
    fmt.Println()
}

func (wsp *WebSocketLoggingProxy) printMessageLog(msgLog MessageLog) {
    direction := "→"
    if msgLog.Direction == "incoming" {
        direction = "←"
    }
    msgType := "Text"
    if msgLog.MessageType == websocket.BinaryMessage {
        msgType = "Binary"
    }
    
    fmt.Println(strings.Repeat("-", 80))
    fmt.Printf("[%s] %s %s Message (%s)\n", 
        msgLog.Timestamp.Format("2006-01-02 15:04:05"), 
        direction, 
        strings.ToUpper(msgLog.Direction),
        msgType)
    fmt.Printf("Remote: %s\n", msgLog.RemoteAddr)
    fmt.Println("\nMessage:")
    fmt.Println(msgLog.Message)
    fmt.Println(strings.Repeat("-", 80))
    fmt.Println()
}

func (wsp *WebSocketLoggingProxy) printFullLog(logEntry LogEntry) {
    fmt.Println(strings.Repeat("=", 80))
    fmt.Println("FULL SESSION LOG")
    fmt.Println(strings.Repeat("=", 80))
    
    // Выводим информацию о подключении
    fmt.Printf("[%s] WebSocket %s\n", 
        logEntry.Connection.Timestamp.Format("2006-01-02 15:04:05"), 
        strings.ToUpper(logEntry.Connection.Event))
    fmt.Printf("Remote: %s\n", logEntry.Connection.RemoteAddr)
    fmt.Printf("URL: %s\n", logEntry.Connection.URL)
    fmt.Printf("Total Messages: %d\n", len(logEntry.Messages))
    
    // Выводим все сообщения
    for i, msg := range logEntry.Messages {
        direction := "→"
        if msg.Direction == "incoming" {
            direction = "←"
        }
        msgType := "Text"
        if msg.MessageType == websocket.BinaryMessage {
            msgType = "Binary"
        }
        
        fmt.Printf("\n[%d] [%s] %s %s (%s)\n", 
            i+1,
            msg.Timestamp.Format("2006-01-02 15:04:05"),
            direction,
            strings.ToUpper(msg.Direction),
            msgType)
        fmt.Println(msg.Message)
    }
    
    fmt.Println(strings.Repeat("=", 80))
    
    // Выводим JSON версию
    if jsonData, err := json.MarshalIndent(logEntry, "", "  "); err == nil {
        fmt.Println("\nJSON Log:")
        fmt.Println(string(jsonData))
        fmt.Println()
    }
}

func main() {
    targetURL := os.Getenv("TARGET_URL")
    if targetURL == "" {
        if len(os.Args) < 2 {
            log.Fatal("Usage: ws-log-proxy <target_ws_url> or set TARGET_URL environment variable")
        }
        targetURL = os.Args[1]
    }
    
    port := os.Getenv("PORT")
    if port == "" {
        port = "8020"
    }
    
    proxy, err := NewWebSocketLoggingProxy(targetURL)
    if err != nil {
        log.Fatalf("Failed to create WebSocket proxy: %v", err)
    }
    
    fmt.Printf("Starting WebSocket log proxy server on port %s\n", port)
    fmt.Printf("Proxying WebSocket connections to: %s\n", targetURL)
    fmt.Println("Press Ctrl+C to stop")
    fmt.Println()
    
    server := &http.Server{
        Addr:         ":" + port,
        Handler:      proxy,
        ReadTimeout:  30 * time.Second,
        WriteTimeout: 30 * time.Second,
    }
    
    log.Fatal(server.ListenAndServe())
}