// Command mitmuxd is the mitmux headless proxy daemon. It owns the
// listening socket and the history database; a TUI or other client
// attaches separately, over a Unix control socket, without interrupting
// capture.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"mitmux/internal/ca"
	"mitmux/internal/ipc"
	"mitmux/internal/proxy"
	"mitmux/internal/store"
)

func main() {
	listen := flag.String("listen", "127.0.0.1:8080", "proxy listen address(es) - comma-separated for more than one, e.g. \"127.0.0.1:8080,127.0.0.1:8081\"")
	caDir := flag.String("ca-dir", "", "directory for CA cert/key and history db (default: XDG config dir)")
	dbPath := flag.String("db", "", "path to history database (default: <ca-dir>/history.db)")
	socketPath := flag.String("socket", "", "control socket path (default: $XDG_RUNTIME_DIR/mitmux.sock, else <ca-dir>/mitmux.sock)")
	installCA := flag.Bool("install-ca", false, "generate the CA if needed, print OS-specific trust-store install steps, and exit (doesn't start the proxy)")
	upstreamProxy := flag.String("upstream-proxy", "", "chain all outbound connections through this HTTP CONNECT proxy (host:port, optional http:// prefix) instead of dialing origins directly")
	flag.Parse()

	var listenAddrs []string
	for _, a := range strings.Split(*listen, ",") {
		if a = strings.TrimSpace(a); a != "" {
			listenAddrs = append(listenAddrs, a)
		}
	}
	if len(listenAddrs) == 0 {
		log.Fatalf("-listen: no addresses given")
	}
	upstream := strings.TrimPrefix(strings.TrimSpace(*upstreamProxy), "http://")

	dir := *caDir
	if dir == "" {
		d, err := ca.Dir()
		if err != nil {
			log.Fatalf("resolve CA dir: %v", err)
		}
		dir = d
	}

	root, err := ca.EnsureCA(dir)
	if err != nil {
		log.Fatalf("load CA: %v", err)
	}

	if *installCA {
		fmt.Printf("CA certificate: %s\n\n", ca.CertPath(dir))
		fmt.Print(ca.InstallInstructions(runtime.GOOS, ca.CertPath(dir)))
		return
	}

	log.Printf("CA ready: %s (install %s/ca.pem in your client's trust store to avoid TLS warnings, or run 'mitmuxd -install-ca' for OS-specific steps)", root.Cert.Subject.CommonName, dir)

	dbFile := *dbPath
	if dbFile == "" {
		dbFile = filepath.Join(dir, "history.db")
	}
	db, err := store.Open(dbFile)
	if err != nil {
		log.Fatalf("open history db: %v", err)
	}
	defer db.Close()
	log.Printf("history db: %s", dbFile)

	sockFile := *socketPath
	if sockFile == "" {
		if rt := os.Getenv("XDG_RUNTIME_DIR"); rt != "" {
			sockFile = filepath.Join(rt, "mitmux.sock")
		} else {
			sockFile = filepath.Join(dir, "mitmux.sock")
		}
	}
	os.Remove(sockFile) // stale socket from an unclean shutdown
	sockLn, err := net.Listen("unix", sockFile)
	if err != nil {
		log.Fatalf("listen on control socket %s: %v", sockFile, err)
	}
	defer os.Remove(sockFile)
	log.Printf("control socket: %s", sockFile)

	srv := proxy.New(listenAddrs, root, db, upstream)
	if upstream != "" {
		log.Printf("chaining outbound connections through upstream proxy %s", upstream)
	}

	hub := ipc.NewHub()
	srv.OnEntry = ipc.LogAndBroadcast(hub)
	ipcSrv := ipc.NewServer(db, hub, srv, strings.Join(listenAddrs, ", "))
	go func() {
		if err := ipcSrv.Serve(sockLn); err != nil {
			log.Printf("control socket: %v", err)
		}
	}()

	errCh := make(chan error, 1)
	go func() {
		errCh <- srv.ListenAndServe()
	}()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)

	select {
	case err := <-errCh:
		if err != nil {
			log.Fatalf("proxy: %v", err)
		}
	case sig := <-sigCh:
		log.Printf("received %s, shutting down", sig)
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := srv.Shutdown(ctx); err != nil {
			log.Fatalf("shutdown: %v", err)
		}
		sockLn.Close()
	}
}