// Package investigation provides bounded, authenticated IP enrichment for the
// operator UI. Passive lookups use PTR and RDAP. Active probing is deliberately
// disabled unless the operator enables it explicitly.
package investigation

import (
	"bytes"
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
)

type Config struct {
	NmapEnabled bool
	NmapBinary  string
	// NmapTimeout bounds one whole scan (OTEL_IP_INVESTIGATION_NMAP_TIMEOUT,
	// a Go duration; default 5m). nmap's own --host-timeout is set below it.
	NmapTimeout time.Duration
}

func ConfigFromEnv() Config {
	binary := os.Getenv("OTEL_IP_INVESTIGATION_NMAP_BINARY")
	if binary == "" {
		binary = "nmap"
	}
	timeout := 5 * time.Minute
	if d, err := time.ParseDuration(os.Getenv("OTEL_IP_INVESTIGATION_NMAP_TIMEOUT")); err == nil && d >= 30*time.Second {
		timeout = d
	}
	return Config{NmapEnabled: strings.EqualFold(os.Getenv("OTEL_IP_INVESTIGATION_NMAP"), "true"), NmapBinary: binary, NmapTimeout: timeout}
}

type Service struct {
	Port     uint16 `json:"port"`
	Protocol string `json:"protocol"`
	State    string `json:"state"`
	Name     string `json:"name,omitempty"`
	Product  string `json:"product,omitempty"`
	Version  string `json:"version,omitempty"`
	// Scripts holds selected NSE output keyed by script id (http-title, ssl-cert).
	Scripts map[string]string `json:"scripts,omitempty"`
}

// Debug describes how an active scan ran, so the UI can explain failures.
type Debug struct {
	Command    string `json:"command"`
	DurationMs int64  `json:"durationMs"`
	ExitError  string `json:"exitError,omitempty"`
	TimedOut   bool   `json:"timedOut,omitempty"`
	Stderr     string `json:"stderr,omitempty"`
	HostState  string `json:"hostState,omitempty"`
	Summary    string `json:"summary,omitempty"`
	Partial    bool   `json:"partial,omitempty"`
}

type errorBody struct {
	Error string `json:"error"`
	Debug *Debug `json:"debug,omitempty"`
}

type PassiveResult struct {
	IP          string   `json:"ip"`
	Hostnames   []string `json:"hostnames,omitempty"`
	RDAP        *RDAP    `json:"rdap,omitempty"`
	Warnings    []string `json:"warnings,omitempty"`
	CompletedAt string   `json:"completedAt"`
}

type RDAP struct {
	Name         string `json:"name,omitempty"`
	Handle       string `json:"handle,omitempty"`
	Start        string `json:"start,omitempty"`
	End          string `json:"end,omitempty"`
	Country      string `json:"country,omitempty"`
	Registry     string `json:"registry,omitempty"`
	Organization string `json:"organization,omitempty"`
	AbuseEmail   string `json:"abuseEmail,omitempty"`
}

type ActiveResult struct {
	IP          string    `json:"ip"`
	Services    []Service `json:"services,omitempty"`
	Debug       *Debug    `json:"debug,omitempty"`
	CompletedAt string    `json:"completedAt"`
}

// Record is one investigation of one address, as returned to the UI and
// persisted. Status is running, done or failed; only finished records are
// stored, so a restart drops in-flight jobs rather than leaving them running.
type Record struct {
	IP         string         `json:"ip"`
	Mode       string         `json:"mode"`
	Status     string         `json:"status"`
	StartedAt  time.Time      `json:"startedAt"`
	FinishedAt *time.Time     `json:"finishedAt,omitempty"`
	Error      string         `json:"error,omitempty"`
	Passive    *PassiveResult `json:"passive,omitempty"`
	Active     *ActiveResult  `json:"active,omitempty"`
	// Debug is kept on failures, where there is no Active result to hold it.
	Debug *Debug `json:"debug,omitempty"`
}

// Store persists finished records. Nil means in-memory only.
type Store interface {
	Save(ctx context.Context, rec Record) error
	Load(ctx context.Context, ip string) ([]Record, error)
}

type Handler struct {
	cfg      Config
	client   *http.Client
	store    Store
	lastScan time.Time
	mu       sync.Mutex
	jobs     map[string]*Record // key: ip|mode; latest in-process record
}

func NewHandler(cfg Config, store Store) *Handler {
	if cfg.NmapTimeout <= 0 {
		cfg.NmapTimeout = 5 * time.Minute
	}
	return &Handler{cfg: cfg, store: store, client: &http.Client{Timeout: 8 * time.Second}, jobs: map[string]*Record{}}
}

func parsePublicIP(r *http.Request) (netip.Addr, bool) {
	ip, err := netip.ParseAddr(strings.TrimSpace(r.PathValue("ip")))
	return ip, err == nil && ip.IsGlobalUnicast() && !ip.IsPrivate() && !ip.IsLoopback()
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "private, no-store")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}

// Start (POST) queues an investigation and returns immediately; poll Get.
func (h *Handler) Start(w http.ResponseWriter, r *http.Request) {
	ip, ok := parsePublicIP(r)
	if !ok {
		writeJSON(w, http.StatusBadRequest, errorBody{Error: "only public unicast IPs can be investigated"})
		return
	}
	var req struct {
		Mode string `json:"mode"`
	}
	if err := json.NewDecoder(io.LimitReader(r.Body, 256)).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, errorBody{Error: "invalid investigation request"})
		return
	}
	if req.Mode != "passive" && req.Mode != "active" {
		writeJSON(w, http.StatusBadRequest, errorBody{Error: "mode must be passive or active"})
		return
	}
	if req.Mode == "active" && !h.cfg.NmapEnabled {
		writeJSON(w, http.StatusNotImplemented, errorBody{Error: "active investigation is disabled; set OTEL_IP_INVESTIGATION_NMAP=true"})
		return
	}
	key := ip.String() + "|" + req.Mode
	h.mu.Lock()
	if cur := h.jobs[key]; cur != nil && cur.Status == "running" {
		rec := *cur
		h.mu.Unlock()
		writeJSON(w, http.StatusAccepted, rec)
		return
	}
	if req.Mode == "active" && time.Since(h.lastScan) < 30*time.Second {
		h.mu.Unlock()
		writeJSON(w, http.StatusTooManyRequests, errorBody{Error: "active investigations are rate limited to one per 30s; retry shortly"})
		return
	}
	if req.Mode == "active" {
		h.lastScan = time.Now()
	}
	rec := &Record{IP: ip.String(), Mode: req.Mode, Status: "running", StartedAt: time.Now().UTC()}
	h.jobs[key] = rec
	snapshot := *rec
	h.mu.Unlock()

	go h.run(key, ip, rec)
	writeJSON(w, http.StatusAccepted, snapshot)
}

// Get returns the latest passive and active record for an address: running
// jobs from memory, otherwise whatever was last stored.
func (h *Handler) Get(w http.ResponseWriter, r *http.Request) {
	ip, ok := parsePublicIP(r)
	if !ok {
		writeJSON(w, http.StatusBadRequest, errorBody{Error: "only public unicast IPs can be investigated"})
		return
	}
	out := map[string]Record{}
	if h.store != nil {
		if recs, err := h.store.Load(r.Context(), ip.String()); err != nil {
			log.Printf("ip investigation: load %s: %v", ip, err)
		} else {
			for _, rec := range recs {
				out[rec.Mode] = rec
			}
		}
	}
	h.mu.Lock()
	for _, mode := range []string{"passive", "active"} {
		if cur := h.jobs[ip.String()+"|"+mode]; cur != nil {
			out[mode] = *cur
		}
	}
	h.mu.Unlock()
	writeJSON(w, http.StatusOK, map[string]any{"ip": ip.String(), "records": out})
}

func (h *Handler) run(key string, ip netip.Addr, rec *Record) {
	timeout := 45 * time.Second
	if rec.Mode == "active" {
		timeout = h.cfg.NmapTimeout
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	var passive *PassiveResult
	var active *ActiveResult
	var err error
	if rec.Mode == "passive" {
		passive = h.passive(ctx, ip)
	} else {
		active, err = h.active(ctx, ip)
	}

	now := time.Now().UTC()
	h.mu.Lock()
	rec.FinishedAt = &now
	rec.Passive, rec.Active = passive, active
	if err != nil {
		rec.Status, rec.Error = "failed", err.Error()
		if active != nil {
			rec.Debug = active.Debug
		}
		log.Printf("ip investigation: nmap %s: %v (debug=%+v)", ip, err, rec.Debug)
	} else {
		rec.Status = "done"
	}
	final := *rec
	h.mu.Unlock()

	if h.store != nil {
		sctx, scancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer scancel()
		if err := h.store.Save(sctx, final); err != nil {
			log.Printf("ip investigation: save %s/%s: %v", ip, final.Mode, err)
		}
	}
}

func (h *Handler) passive(ctx context.Context, ip netip.Addr) *PassiveResult {
	result := &PassiveResult{IP: ip.String(), CompletedAt: time.Now().UTC().Format(time.RFC3339)}
	if names, err := net.DefaultResolver.LookupAddr(ctx, ip.String()); err == nil {
		for _, name := range names {
			result.Hostnames = append(result.Hostnames, strings.TrimSuffix(name, "."))
		}
	} else {
		result.Warnings = append(result.Warnings, "reverse DNS lookup returned no name")
	}
	if rdap, err := h.lookupRDAP(ctx, ip); err == nil {
		result.RDAP = rdap
	} else {
		result.Warnings = append(result.Warnings, "RDAP lookup unavailable")
	}
	return result
}

// IANA's RDAP bootstrap format is an array of tuples, not an array of
// objects: [ ["prefix", ...], ["https://rdap.example/", ...] ].
type bootstrap struct {
	Services [][]json.RawMessage `json:"services"`
}

func (b bootstrap) serviceFor(ip netip.Addr) (string, bool) {
	for _, service := range b.Services {
		if len(service) < 2 {
			continue
		}
		var prefixes, urls []string
		if json.Unmarshal(service[0], &prefixes) != nil || json.Unmarshal(service[1], &urls) != nil || len(urls) == 0 {
			continue
		}
		for _, prefix := range prefixes {
			p, err := netip.ParsePrefix(prefix)
			if err == nil && p.Contains(ip) {
				return urls[0], true
			}
		}
	}
	return "", false
}

type rdapEntity struct {
	Roles    []string     `json:"roles"`
	VCard    []any        `json:"vcardArray"`
	Entities []rdapEntity `json:"entities"`
}

type rdapResponse struct {
	Name     string       `json:"name"`
	Handle   string       `json:"handle"`
	Start    string       `json:"startAddress"`
	End      string       `json:"endAddress"`
	Country  string       `json:"country"`
	Entities []rdapEntity `json:"entities"`
}

func (h *Handler) lookupRDAP(ctx context.Context, ip netip.Addr) (*RDAP, error) {
	bootstrapURL := "https://data.iana.org/rdap/ipv4.json"
	if ip.Is6() {
		bootstrapURL = "https://data.iana.org/rdap/ipv6.json"
	}
	var b bootstrap
	if err := h.getJSON(ctx, bootstrapURL, &b); err != nil {
		return nil, err
	}
	serviceURL, ok := b.serviceFor(ip)
	if !ok {
		return nil, fmt.Errorf("no RDAP service for %s", ip)
	}
	var out rdapResponse
	if err := h.getJSON(ctx, strings.TrimRight(serviceURL, "/")+"/ip/"+ip.String(), &out); err != nil {
		return nil, err
	}
	organization, abuseEmail := rdapContacts(out.Entities)
	return &RDAP{
		Name: out.Name, Handle: out.Handle, Start: out.Start, End: out.End,
		Country: out.Country, Registry: serviceURL, Organization: organization, AbuseEmail: abuseEmail,
	}, nil

}

func rdapContacts(entities []rdapEntity) (organization, abuseEmail string) {
	for _, entity := range entities {
		fn, email := vcardFields(entity.VCard)
		for _, role := range entity.Roles {
			switch role {
			case "registrant", "administrative", "technical":
				if organization == "" {
					organization = fn
				}
			case "abuse":
				if abuseEmail == "" {
					abuseEmail = email
				}
			}
		}
		if childOrg, childEmail := rdapContacts(entity.Entities); organization == "" {
			organization = childOrg
			if abuseEmail == "" {
				abuseEmail = childEmail
			}
		} else if abuseEmail == "" {
			_, abuseEmail = rdapContacts(entity.Entities)
		}
	}
	return organization, abuseEmail
}

func (h *Handler) getJSON(ctx context.Context, url string, dst any) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/rdap+json, application/json")
	resp, err := h.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("HTTP %s", resp.Status)
	}
	return json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(dst)
}

// vcardFields extracts fn and email from a jCard ["vcard", [[name, params, type, value], ...]].
func vcardFields(card []any) (fn, email string) {
	if len(card) < 2 {
		return
	}
	props, _ := card[1].([]any)
	for _, p := range props {
		prop, _ := p.([]any)
		if len(prop) < 4 {
			continue
		}
		name, _ := prop[0].(string)
		val, _ := prop[3].(string)
		switch name {
		case "fn":
			fn = val
		case "email":
			if email == "" {
				email = val
			}
		}
	}
	return
}

type nmapXML struct {
	Hosts []struct {
		Status struct {
			State  string `xml:"state,attr"`
			Reason string `xml:"reason,attr"`
		} `xml:"status"`
		Ports struct {
			ExtraPorts []struct {
				State string `xml:"state,attr"`
				Count int    `xml:"count,attr"`
			} `xml:"extraports"`
			Ports []struct {
				Protocol string `xml:"protocol,attr"`
				Port     uint16 `xml:"portid,attr"`
				State    struct {
					Value string `xml:"state,attr"`
				} `xml:"state"`
				Service struct {
					Name    string `xml:"name,attr"`
					Product string `xml:"product,attr"`
					Version string `xml:"version,attr"`
				} `xml:"service"`
				Scripts []struct {
					ID     string `xml:"id,attr"`
					Output string `xml:"output,attr"`
				} `xml:"script"`
			} `xml:"port"`
		} `xml:"ports"`
	} `xml:"host"`
	Finished struct {
		Summary string `xml:"summary,attr"`
	} `xml:"runstats>finished"`
}

func (h *Handler) active(ctx context.Context, ip netip.Addr) (*ActiveResult, error) {
	// --host-timeout makes nmap stop itself and emit valid XML with whatever it
	// found; the context deadline is only a backstop and stays well above it.
	hostTimeout := fmt.Sprintf("%ds", int((h.cfg.NmapTimeout - 15*time.Second).Seconds()))
	args := []string{"-Pn", "-n", "-T4", "--max-retries", "1", "--host-timeout", hostTimeout, "--version-light", "-sV",
		"--script", "ssl-cert,http-title", "--top-ports", "100", "--open", "-oX", "-", ip.String()}
	result := &ActiveResult{IP: ip.String(), Debug: &Debug{Command: h.cfg.NmapBinary + " " + strings.Join(args, " ")}}
	start := time.Now()
	cmd := exec.CommandContext(ctx, h.cfg.NmapBinary, args...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	err := cmd.Run()
	result.Debug.DurationMs = time.Since(start).Milliseconds()
	result.Debug.Stderr = truncate(strings.TrimSpace(stderr.String()), 2000)
	result.CompletedAt = time.Now().UTC().Format(time.RFC3339)
	if err != nil {
		result.Debug.ExitError = err.Error()
		result.Debug.TimedOut = ctx.Err() != nil
		if result.Debug.TimedOut {
			return result, fmt.Errorf("nmap exceeded the %s scan deadline and was killed", h.cfg.NmapTimeout)
		}
		return result, fmt.Errorf("nmap failed: %w", err)
	}
	var parsed nmapXML
	if err := xml.Unmarshal(stdout.Bytes(), &parsed); err != nil {
		result.Debug.ExitError = "unparseable nmap XML: " + err.Error()
		return result, fmt.Errorf("could not parse nmap output: %w", err)
	}
	result.Debug.Summary = parsed.Finished.Summary
	for _, host := range parsed.Hosts {
		result.Debug.HostState = host.Status.State
		if host.Status.Reason != "" {
			result.Debug.HostState += " (" + host.Status.Reason + ")"
		}
		for _, x := range host.Ports.ExtraPorts {
			result.Debug.Summary += fmt.Sprintf("; %d ports %s", x.Count, x.State)
		}
		for _, p := range host.Ports.Ports {
			svc := Service{Port: p.Port, Protocol: p.Protocol, State: p.State.Value, Name: p.Service.Name, Product: p.Service.Product, Version: p.Service.Version}
			for _, sc := range p.Scripts {
				if svc.Scripts == nil {
					svc.Scripts = map[string]string{}
				}
				svc.Scripts[sc.ID] = truncate(sc.Output, 600)
			}
			result.Services = append(result.Services, svc)
		}
	}
	if strings.Contains(result.Debug.Summary, "timed out") || len(parsed.Hosts) == 0 {
		result.Debug.Partial = true
	}
	return result, nil
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}