package main

import (
    "fmt"
    "net/http"
    "os"
    "path/filepath"
    "sort"
    "strings"
    "sync"
    "time"

    "acme-reverseproxy/config"
    "acme-reverseproxy/nfs"
    "acme-reverseproxy/proxymap"
    "acme-reverseproxy/streams"

    "github.com/BurntSushi/toml"
    "github.com/sirupsen/logrus"
    "github.com/urfave/cli"
    "golang.org/x/crypto/acme/autocert"
)

var cfg config.Config

func main() {
    app := cli.NewApp()
    app.Name = "acme-reverseproxy"
    app.Usage = "A TLS-serving reverse-proxy, with the certificates generated from LetsEncrypt"
    app.Authors = []cli.Author{{Name: "Forc", Email: "redirsvr@mail.ru"}}
    app.Flags = []cli.Flag{cli.BoolFlag{Name: "debug,D", Usage: "debug output"}}
    app.Commands = []cli.Command{
        {
            Name:        "gen",
            Description: "generators of sorts",
            Subcommands: []cli.Command{
                {
                    Name:        "config",
                    Description: "generate a sample mapping configuration",
                    Action:      GenConfigAction,
                },
            },
        },
        {
            Name:        "srv",
            Description: "Start the reverseproxy server",
            Action:      SrvCommand,
            Flags: []cli.Flag{
                cli.StringFlag{
                    Name:  "config",
                    Value: filepath.Join(os.Getenv("HOME"), ".acme-reverseproxy.toml"),
                    Usage: "Configuration of mapping of hostname -> listener",
                },
            },
            Before: BeforeAction,
        },
    }
    sort.Sort(cli.FlagsByName(app.Flags))
    sort.Sort(cli.CommandsByName(app.Commands))
    app.Run(os.Args)
}

func BeforeAction(c *cli.Context) error {
    if c.GlobalBool("debug") {
        logrus.SetLevel(logrus.DebugLevel)
        logrus.Info("Debug mode enabled")
    } else {
        logrus.SetLevel(logrus.InfoLevel)
    }
    logrus.Infof("Loading config from: %s", c.String("config"))
    buf, err := os.ReadFile(c.String("config"))
    if err != nil {
        if os.IsNotExist(err) {
            logrus.Errorf("No config file found at %q. Try 'gen config'", c.String("config"))
        }
        return err
    }
    tmpConfig := config.Config{}
    if err := toml.Unmarshal(buf, &tmpConfig); err != nil {
        logrus.Errorf("Failed to parse config: %v", err)
        return err
    }
    cfg = tmpConfig
    logrus.Infof("Config loaded: CA.Email=%s, CA.CacheDir=%s, WellKnownDir=%s", 
        cfg.CA.Email, cfg.CA.CacheDir, cfg.WellKnownDir)
    logrus.Infof("Configured mappings: %d domains", len(cfg.Mapping))
    for host, target := range cfg.Mapping {
        logrus.Infof("  %s -> %s", host, target)
    }
    logrus.Infof("Configured streams: %d", len(cfg.Streams))
    for domain, stream := range cfg.Streams {
        logrus.Infof("  %s -> %s (%s)", domain, stream.Target, stream.Protocol)
    }
    return nil
}

func GenConfigAction(c *cli.Context) error {
    tmpConfig := config.Config{
        CA: config.CA{
            Email:    "admin@p42.ru",
            CacheDir: "/tmp/acme-reverseproxy",
        },
        Mapping: map[string]string{
            "p42.ru": "http://localhost:5000",
        },
        WellKnownDir: "/tmp/.well-known",
        NFS: config.NFSConfig{
            Enabled:    false,
            Server:     "192.168.1.100",
            ExportPath: "/export/acme",
            MountPoint: "/mnt/acme-wellknown",
            Options:    "rw,vers=4.1,timeo=50,retrans=2",
        },
    }
    e := toml.NewEncoder(os.Stdout)
    if err := e.Encode(tmpConfig); err != nil {
        return err
    }
    return nil
}

func SrvCommand(c *cli.Context) error {
    // Инициализируем NFS менеджер
    nfs.Init(nfs.Config{
        Enabled:    cfg.NFS.Enabled,
        Server:     cfg.NFS.Server,
        ExportPath: cfg.NFS.ExportPath,
        MountPoint: cfg.NFS.MountPoint,
        Options:    cfg.NFS.Options,
    })

    // Если NFS включен, монтируем его при старте
    if cfg.NFS.Enabled {
        logrus.Info("Initializing NFS mount at startup...")
        if err := nfs.EnsureMounted(); err != nil {
            logrus.Warnf("NFS mount failed at startup (will retry on demand): %v", err)
        } else {
            logrus.Info("NFS mounted successfully at startup")
        }
    }

    // Инициализируем менеджер стримов
    streamManager := streams.NewStreamManager(cfg.Streams)
    if err := streamManager.Start(); err != nil {
        logrus.Errorf("Failed to start stream manager: %v", err)
        return cli.NewExitError(err, 2)
    }

    list := []string{}
    for key := range cfg.Mapping {
        if key != "" {
            list = append(list, key)
        }
    }
    if len(list) == 0 {
        logrus.Error("No domains configured in Mapping section")
        return cli.NewExitError("No domains configured", 2)
    }
    logrus.Infof("Initializing reverse proxy for %d domains: %s", len(list), strings.Join(list, ", "))
    rpm, err := proxymap.ToReverseProxyMap(cfg.Mapping)
    if err != nil {
        logrus.Errorf("Failed to create reverse proxy map: %v", err)
        return cli.NewExitError(err, 2)
    }
    wellKnownDir := cfg.WellKnownDir
    if wellKnownDir == "" {
        wellKnownDir = "/tmp/.well-known"
    }
    // Если NFS включен и указан MountPoint, используем его для .well-known
    if cfg.NFS.Enabled && cfg.NFS.MountPoint != "" {
        wellKnownDir = cfg.NFS.MountPoint
        logrus.Infof("Using NFS mount point for .well-known: %s", wellKnownDir)
    }
    logrus.Infof("Well-known directory: %s", wellKnownDir)
    rph := proxymap.NewReverseProxiesHandlerWithWellKnown(rpm, wellKnownDir)
    logrus.Infof("ACME whitelist: %s", strings.Join(list, ", "))
    m := autocert.Manager{
        Prompt:     autocert.AcceptTOS,
        HostPolicy: autocert.HostWhitelist(list...),
    }
    if cfg.CA.Email != "" {
        m.Email = cfg.CA.Email
        logrus.Infof("ACME email: %s", cfg.CA.Email)
    }
    if cfg.CA.CacheDir != "" {
        m.Cache = autocert.DirCache(cfg.CA.CacheDir)
        logrus.Infof("ACME cache directory: %s", cfg.CA.CacheDir)
    }
    httpsListener := m.Listener()
    httpsServer := &http.Server{
        Handler: rph,
    }
    acmeHTTPHandler := m.HTTPHandler(rph)
    httpHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if proxymap.IsAcmeWellKnownPath(r.URL.Path) {
            rph.ServeHTTP(w, r)
            return
        }
        acmeHTTPHandler.ServeHTTP(w, r)
    })
    httpServer := &http.Server{Addr: ":80", Handler: httpHandler}
    logrus.Info(strings.Repeat("=", 61))
    logrus.Info("Starting reverse proxy servers...")
    logrus.Info(strings.Repeat("=", 61))
    var wg sync.WaitGroup
    var firstError error
    var mu sync.Mutex
    
    // Запускаем HTTP/HTTPS серверы
    wg.Add(2)
    go func() {
        defer wg.Done()
        logrus.Info("Starting HTTPS server on port 443...")
        if err := httpsServer.Serve(httpsListener); err != nil {
            mu.Lock()
            if firstError == nil {
                firstError = err
            }
            mu.Unlock()
            logrus.Errorf("HTTPS server stopped with error: %v", err)
        } else {
            logrus.Info("HTTPS server stopped normally")
        }
    }()
    go func() {
        defer wg.Done()
        logrus.Info("Starting HTTP server on port 80 (redirecting to HTTPS)...")
        if err := httpServer.ListenAndServe(); err != nil {
            mu.Lock()
            if firstError == nil {
                firstError = err
            }
            mu.Unlock()
            logrus.Errorf("HTTP server stopped with error: %v", err)
        } else {
            logrus.Info("HTTP server stopped normally")
        }
    }()
    time.Sleep(100 * time.Millisecond)
    logrus.Info("Servers started successfully. Waiting for requests...")
    logrus.Info("Press Ctrl+C to stop")
    wg.Wait()
    
    // При остановке останавливаем стримы
    streamManager.Stop()
    
    // При остановке размонтируем NFS если он был смонтирован
    if cfg.NFS.Enabled {
        if err := nfs.DefaultManager.Unmount(); err != nil {
            logrus.Warnf("Failed to unmount NFS on shutdown: %v", err)
        }
    }
    if firstError != nil {
        logrus.Errorf("Server exited with error: %v", firstError)
        return cli.NewExitError(firstError, 2)
    }
    logrus.Info("All servers stopped")
    return nil
}

func MakeDir(path string) string {
    err := os.MkdirAll(path, 0777)
    if err != nil {
        fmt.Println(err, path)
        return path
    }
    return path
}
