package main

import (
    "bytes"
    "encoding/json"
    "fmt"
    "io"
    "log"
    "net/http"
    "net/http/httputil"
    "net/url"
    "os"
    "strings"
    "time"
)

type LoggingProxy struct {
    target *url.URL
    proxy  *httputil.ReverseProxy
}

type RequestLog struct {
    Timestamp  time.Time              `json:"timestamp"`
    Method     string                 `json:"method"`
    URL        string                 `json:"url"`
    FullURL    string                 `json:"full_url"`
    TargetURL  string                 `json:"target_url"`
    Headers    map[string][]string    `json:"headers"`
    Body       string                 `json:"body"`
    RemoteAddr string                 `json:"remote_addr"`
}

type ResponseLog struct {
    Timestamp  time.Time          `json:"timestamp"`
    StatusCode int                `json:"status_code"`
    Headers    map[string][]string `json:"headers"`
    Body       string             `json:"body"`
    Duration   string             `json:"duration"`
}

type LogEntry struct {
    Request  RequestLog  `json:"request"`
    Response ResponseLog `json:"response"`
}

func NewLoggingProxy(targetURL string) (*LoggingProxy, error) {
    target, err := url.Parse(targetURL)
    if err != nil {
        return nil, fmt.Errorf("invalid target URL: %w", err)
    }
    proxy := httputil.NewSingleHostReverseProxy(target)
    originalDirector := proxy.Director
    proxy.Director = func(req *http.Request) {
        originalDirector(req)
        // Устанавливаем правильный хост
        req.Host = target.Host
        // Убеждаемся, что схема правильная
        req.URL.Scheme = target.Scheme
        req.URL.Host = target.Host
        // Сохраняем оригинальный путь и query параметры
        // (они уже должны быть сохранены через originalDirector)
    }
    return &LoggingProxy{
        target: target,
        proxy:  proxy,
    }, nil
}

func (lp *LoggingProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    start := time.Now()
    
    // Логируем запрос перед отправкой
    reqLog := lp.logRequest(r)
    
    // Создаем recorder для перехвата ответа
    recorder := &responseRecorder{
        ResponseWriter: w,
        statusCode:     http.StatusOK,
        body:           &bytes.Buffer{},
        headers:        make(http.Header),
    }
    
    // Настраиваем ModifyResponse для перехвата метаданных ответа
    originalModifyResponse := lp.proxy.ModifyResponse
    lp.proxy.ModifyResponse = func(resp *http.Response) error {
        if originalModifyResponse != nil {
            if err := originalModifyResponse(resp); err != nil {
                return err
            }
        }
        // Сохраняем статус код из ответа для логирования
        recorder.statusCode = resp.StatusCode
        recorder.statusSetFromResp = true
        
        // Сохраняем заголовки для логирования
        // ReverseProxy автоматически скопирует заголовки в ResponseWriter
        if recorder.headers == nil {
            recorder.headers = make(http.Header)
        }
        for k, v := range resp.Header {
            recorder.headers[k] = make([]string, len(v))
            copy(recorder.headers[k], v)
        }
        
        // Важно: не трогаем resp.Body, чтобы ReverseProxy мог правильно передать его клиенту
        return nil
    }
    
    // Выполняем проксирование запроса
    lp.proxy.ServeHTTP(recorder, r)
    
    // Логируем ответ после получения
    duration := time.Since(start)
    respLog := ResponseLog{
        Timestamp:  time.Now(),
        StatusCode: recorder.statusCode,
        Headers:    recorder.headers,
        Body:       recorder.body.String(),
        Duration:   duration.String(),
    }
    logEntry := LogEntry{
        Request:  reqLog,
        Response: respLog,
    }
    lp.printLog(logEntry)
}

func (lp *LoggingProxy) logRequest(r *http.Request) RequestLog {
    // Читаем тело запроса
    bodyBytes, _ := io.ReadAll(r.Body)
    r.Body.Close()
    // Восстанавливаем тело для проксирования
    r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
    
    // Форматируем тело для логирования
    bodyStr := string(bodyBytes)
    if len(bodyStr) > 0 {
        var jsonData interface{}
        if err := json.Unmarshal(bodyBytes, &jsonData); err == nil {
            if formatted, err := json.MarshalIndent(jsonData, "", "  "); err == nil {
                bodyStr = string(formatted)
            }
        }
    }
    
    // Формируем полный URL запроса
    fullURL := r.URL.String()
    if r.URL.Scheme == "" {
        scheme := "http"
        if r.TLS != nil {
            scheme = "https"
        }
        fullURL = fmt.Sprintf("%s://%s%s", scheme, r.Host, r.URL.String())
    }
    
    // Формируем целевой URL
    targetURL := fmt.Sprintf("%s://%s%s", lp.target.Scheme, lp.target.Host, r.URL.Path)
    if r.URL.RawQuery != "" {
        targetURL += "?" + r.URL.RawQuery
    }
    
    return RequestLog{
        Timestamp:  time.Now(),
        Method:     r.Method,
        URL:        r.URL.String(),
        FullURL:    fullURL,
        TargetURL:  targetURL,
        Headers:    r.Header,
        Body:       bodyStr,
        RemoteAddr: r.RemoteAddr,
    }
}

func (lp *LoggingProxy) printLog(entry LogEntry) {
    fmt.Println(strings.Repeat("=", 80))
    fmt.Printf("[%s] %s %s\n", entry.Request.Timestamp.Format("2006-01-02 15:04:05"), entry.Request.Method, entry.Request.URL)
    fmt.Printf("From URL: %s\n", entry.Request.FullURL)
    fmt.Printf("To URL: %s\n", entry.Request.TargetURL)
    fmt.Printf("Remote: %s\n", entry.Request.RemoteAddr)
    if len(entry.Request.Headers) > 0 {
        fmt.Println("\nRequest Headers:")
        for k, v := range entry.Request.Headers {
            fmt.Printf("  %s: %s\n", k, strings.Join(v, ", "))
        }
    }
    if entry.Request.Body != "" {
        fmt.Println("\nRequest Body:")
        fmt.Println(entry.Request.Body)
    }
    fmt.Printf("\n[%s] Response: %d (%s)\n", entry.Response.Timestamp.Format("2006-01-02 15:04:05"), entry.Response.StatusCode, entry.Response.Duration)
    if len(entry.Response.Headers) > 0 {
        fmt.Println("\nResponse Headers:")
        for k, v := range entry.Response.Headers {
            fmt.Printf("  %s: %s\n", k, strings.Join(v, ", "))
        }
    }
    if entry.Response.Body != "" {
        fmt.Println("\nResponse Body:")
        fmt.Println(entry.Response.Body)
    }
    fmt.Println(strings.Repeat("=", 80))
    fmt.Println()
    if jsonData, err := json.MarshalIndent(entry, "", "  "); err == nil {
        fmt.Println("JSON Log:")
        fmt.Println(string(jsonData))
        fmt.Println()
    }
}

type responseRecorder struct {
    http.ResponseWriter
    statusCode    int
    body          *bytes.Buffer
    headers       http.Header
    wroteHeader   bool
    statusSetFromResp bool
}

func (rr *responseRecorder) Header() http.Header {
    return rr.ResponseWriter.Header()
}

func (rr *responseRecorder) WriteHeader(code int) {
    if !rr.wroteHeader {
        // Используем статус код из ModifyResponse, если он был установлен, иначе используем переданный
        if !rr.statusSetFromResp {
            rr.statusCode = code
        }
        rr.wroteHeader = true
        
        // Сохраняем все заголовки перед отправкой (для логирования)
        if rr.headers == nil {
            rr.headers = make(http.Header)
        }
        // Обновляем заголовки из ResponseWriter (они могут быть изменены ReverseProxy)
        for k, v := range rr.ResponseWriter.Header() {
            rr.headers[k] = make([]string, len(v))
            copy(rr.headers[k], v)
        }
        // Вызываем WriteHeader оригинального ResponseWriter с правильным статус кодом
        rr.ResponseWriter.WriteHeader(rr.statusCode)
    }
}

func (rr *responseRecorder) Write(b []byte) (int, error) {
    // Если WriteHeader еще не вызван, вызываем его со статусом 200
    if !rr.wroteHeader {
        rr.WriteHeader(http.StatusOK)
    }
    
    // Сохраняем тело для логирования
    if rr.body == nil {
        rr.body = &bytes.Buffer{}
    }
    rr.body.Write(b)
    
    // Важно: записываем в оригинальный ResponseWriter, чтобы ответ вернулся клиенту
    return rr.ResponseWriter.Write(b)
}

func main() {
    targetURL := os.Getenv("TARGET_URL")
    if targetURL == "" {
        if len(os.Args) < 2 {
            log.Fatal("Usage: log-proxy <target_url> or set TARGET_URL environment variable")
        }
        targetURL = os.Args[1]
    }
    port := os.Getenv("PORT")
    if port == "" {
        port = "8020"
    }
    proxy, err := NewLoggingProxy(targetURL)
    if err != nil {
        log.Fatalf("Failed to create proxy: %v", err)
    }
    fmt.Printf("Starting log proxy server on port %s\n", port)
    fmt.Printf("Proxying requests 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())
}