package dig

import (
	"context"
	"fmt"
	"os/exec"
	"strings"

	"github.com/anotherhadi/iknowyou/internal/tools"
)

const (
	name        = "dig"
	description = "DNS lookup querying A, AAAA, MX, NS, TXT, and SOA records for a domain, or reverse DNS (PTR) for an IP."
	link        = "https://linux.die.net/man/1/dig"
	icon        = ""
)

var recordTypes = []string{"A", "AAAA", "MX", "NS", "TXT", "SOA"}

type Runner struct{}

func New() tools.ToolRunner { return &Runner{} }

func (r *Runner) Name() string        { return name }
func (r *Runner) Description() string { return description }
func (r *Runner) Link() string        { return link }
func (r *Runner) Icon() string        { return icon }

func (r *Runner) InputTypes() []tools.InputType {
	return []tools.InputType{tools.InputTypeDomain, tools.InputTypeIP}
}

func (r *Runner) Available() (bool, string) {
	if _, err := exec.LookPath("dig"); err != nil {
		return false, "dig binary not found in PATH"
	}
	return true, ""
}

func (r *Runner) Dependencies() []string { return []string{"dig"} }

func (r *Runner) Run(ctx context.Context, target string, inputType tools.InputType, out chan<- tools.Event) error {
	defer close(out)

	var sb strings.Builder
	totalCount := 0

	if inputType == tools.InputTypeIP {
		cmd := exec.CommandContext(ctx, "dig", "-x", target, "+noall", "+answer")
		output, err := cmd.Output()
		if err != nil && ctx.Err() != nil {
			out <- tools.Event{Tool: name, Type: tools.EventTypeError, Payload: "scan cancelled"}
			out <- tools.Event{Tool: name, Type: tools.EventTypeCount, Payload: 0}
			out <- tools.Event{Tool: name, Type: tools.EventTypeDone}
			return nil
		}
		result := strings.TrimSpace(string(output))
		if result != "" {
			sb.WriteString("=== Reverse DNS (PTR) ===\n")
			sb.WriteString(result)
			totalCount += strings.Count(result, "\n") + 1
		}
	} else {
		for _, rtype := range recordTypes {
			if ctx.Err() != nil {
				break
			}
			cmd := exec.CommandContext(ctx, "dig", target, rtype, "+noall", "+answer")
			output, _ := cmd.Output()
			result := strings.TrimSpace(string(output))
			if result == "" {
				continue
			}
			sb.WriteString(fmt.Sprintf("=== %s ===\n", rtype))
			sb.WriteString(result)
			sb.WriteString("\n\n")
			totalCount += strings.Count(result, "\n") + 1
		}
	}

	if ctx.Err() != nil {
		out <- tools.Event{Tool: name, Type: tools.EventTypeError, Payload: "scan cancelled"}
	} else if sb.Len() > 0 {
		out <- tools.Event{Tool: name, Type: tools.EventTypeOutput, Payload: strings.TrimSpace(sb.String())}
	}
	out <- tools.Event{Tool: name, Type: tools.EventTypeCount, Payload: totalCount}
	out <- tools.Event{Tool: name, Type: tools.EventTypeDone}
	return nil
}