package core

import (
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/Yenn503/NecropolisC2/pkg/cryptography"
)

// colorLogWriter wraps a writer and colorizes known log prefixes.
type colorLogWriter struct{ w io.Writer }

func (c *colorLogWriter) Write(p []byte) (int, error) {
	s := string(p)
	switch {
	case strings.Contains(s, "[operator] "):
		s = strings.Replace(s, "[operator] ", boldCyan("[operator] "), 1)
	case strings.Contains(s, "[implant] "):
		s = strings.Replace(s, "[implant] ", boldGreen("[implant] "), 1)
	case strings.Contains(s, "[evasion] "):
		s = strings.Replace(s, "[evasion] ", magentaBold("[evasion] "), 1)
	case strings.Contains(s, "[socks] "):
		s = strings.Replace(s, "[socks] ", dimYellow("[socks] "), 1)
	case strings.Contains(s, "[node] "):
		s = strings.Replace(s, "[node] ", dim+"[node]"+reset+" ", 1)
	case strings.Contains(s, "[WARN] "):
		s = strings.Replace(s, "[WARN]", boldRed("[WARN]"), 1)
	case strings.Contains(s, "DISCONNECTED"):
		s = strings.Replace(s, "DISCONNECTED", boldRed("DISCONNECTED"), 1)
	}
	return c.w.Write([]byte(s))
}

func Run(relayAddrs []string) error {
	log.SetOutput(&colorLogWriter{w: os.Stderr})
	if err := os.MkdirAll(necropolisDir(), 0700); err != nil {
		return fmt.Errorf("create necropolis directory: %w", err)
	}

	keys, err := cryptography.LoadOrGenerateOperatorKey(keyPath())
	if err != nil {
		return fmt.Errorf("load operator key: %w", err)
	}

	pubPath := pubKeyPath()
	if _, err := os.Stat(pubPath); os.IsNotExist(err) {
		pubBytes, err := crypto.MarshalPublicKey(keys.PublicKey)
		if err == nil {
			if err := os.WriteFile(pubPath, pubBytes, 0644); err == nil {
				log.Printf("[operator] public key exported: %s", pubPath)
			}
		}
	}

	log.Printf("[operator] peer ID: %s", keys.PeerID.String())

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	op, err := NewOperator(ctx, keys, relayAddrs)
	if err != nil {
		return fmt.Errorf("create operator: %w", err)
	}
	defer op.Close()

	if err := op.Start(); err != nil {
		return fmt.Errorf("start operator: %w", err)
	}

	op.RunCLI()
	return nil
}

func necropolisDir() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return ".necropolis"
	}
	return filepath.Join(home, ".necropolis")
}

func keyPath() string {
	return filepath.Join(necropolisDir(), "operator.key")
}

func pubKeyPath() string {
	return filepath.Join(necropolisDir(), "operator.pub")
}