package sources

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"golang.org/x/time/rate"

	"github.com/salvacybersec/keyhunter/pkg/providers"
	"github.com/salvacybersec/keyhunter/pkg/recon"
)

// SecurityTrailsSource searches SecurityTrails DNS/subdomain data for API key
// exposure. It enumerates subdomains for a target domain and probes config
// endpoints, and also checks DNS history records (TXT records may contain keys).
type SecurityTrailsSource struct {
	APIKey   string
	BaseURL  string
	Registry *providers.Registry
	Limiters *recon.LimiterRegistry
	Client   *Client

	// ProbeBaseURL overrides the scheme+host used when probing discovered
	// subdomains. Tests set this to the httptest server URL.
	ProbeBaseURL string
}

var _ recon.ReconSource = (*SecurityTrailsSource)(nil)

func (s *SecurityTrailsSource) Name() string         { return "securitytrails" }
func (s *SecurityTrailsSource) RateLimit() rate.Limit { return rate.Every(2 * time.Second) }
func (s *SecurityTrailsSource) Burst() int            { return 5 }
func (s *SecurityTrailsSource) RespectsRobots() bool  { return false }

func (s *SecurityTrailsSource) Enabled(_ recon.Config) bool {
	return s.APIKey != ""
}

// securityTrailsSubdomains represents the subdomain listing API response.
type securityTrailsSubdomains struct {
	Subdomains []string `json:"subdomains"`
}

func (s *SecurityTrailsSource) Sweep(ctx context.Context, query string, out chan<- recon.Finding) error {
	base := s.BaseURL
	if base == "" {
		base = "https://api.securitytrails.com/v1"
	}
	client := s.Client
	if client == nil {
		client = NewClient()
	}

	if query == "" || !strings.Contains(query, ".") {
		return nil
	}

	// Phase 1: Enumerate subdomains.
	if s.Limiters != nil {
		if err := s.Limiters.Wait(ctx, s.Name(), s.RateLimit(), s.Burst(), false); err != nil {
			return err
		}
	}

	subURL := fmt.Sprintf("%s/domain/%s/subdomains?children_only=false", base, query)
	subReq, err := http.NewRequestWithContext(ctx, http.MethodGet, subURL, nil)
	if err != nil {
		return err
	}
	subReq.Header.Set("APIKEY", s.APIKey)

	subResp, err := client.Do(ctx, subReq)
	if err != nil {
		return nil // non-fatal
	}

	subData, err := io.ReadAll(io.LimitReader(subResp.Body, 512*1024))
	_ = subResp.Body.Close()
	if err != nil {
		return nil
	}

	var subResult securityTrailsSubdomains
	if err := json.Unmarshal(subData, &subResult); err != nil {
		return nil
	}

	// Build FQDNs and limit to 20.
	var fqdns []string
	for _, sub := range subResult.Subdomains {
		fqdns = append(fqdns, sub+"."+query)
		if len(fqdns) >= 20 {
			break
		}
	}

	// Probe config endpoints on each subdomain.
	probeClient := &http.Client{Timeout: 5 * time.Second}
	for _, fqdn := range fqdns {
		if err := ctx.Err(); err != nil {
			return err
		}
		s.probeSubdomain(ctx, probeClient, fqdn, out)
	}

	// Phase 2: Check DNS history for key patterns in TXT records.
	if s.Limiters != nil {
		if err := s.Limiters.Wait(ctx, s.Name(), s.RateLimit(), s.Burst(), false); err != nil {
			return err
		}
	}

	dnsURL := fmt.Sprintf("%s/domain/%s", base, query)
	dnsReq, err := http.NewRequestWithContext(ctx, http.MethodGet, dnsURL, nil)
	if err != nil {
		return nil
	}
	dnsReq.Header.Set("APIKEY", s.APIKey)

	dnsResp, err := client.Do(ctx, dnsReq)
	if err != nil {
		return nil
	}

	dnsData, err := io.ReadAll(io.LimitReader(dnsResp.Body, 512*1024))
	_ = dnsResp.Body.Close()
	if err != nil {
		return nil
	}

	if ciLogKeyPattern.Match(dnsData) {
		out <- recon.Finding{
			ProviderName: query,
			Source:       dnsURL,
			SourceType:   "recon:securitytrails",
			Confidence:   "medium",
			DetectedAt:   time.Now(),
		}
	}

	return nil
}

// probeSubdomain checks well-known config endpoints for key patterns.
func (s *SecurityTrailsSource) probeSubdomain(ctx context.Context, probeClient *http.Client, subdomain string, out chan<- recon.Finding) {
	for _, ep := range configProbeEndpoints {
		if err := ctx.Err(); err != nil {
			return
		}

		var probeURL string
		if s.ProbeBaseURL != "" {
			probeURL = s.ProbeBaseURL + "/" + subdomain + ep
		} else {
			probeURL = "https://" + subdomain + ep
		}

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, probeURL, nil)
		if err != nil {
			continue
		}

		resp, err := probeClient.Do(req)
		if err != nil {
			continue
		}

		body, err := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
		_ = resp.Body.Close()
		if err != nil {
			continue
		}

		if resp.StatusCode == http.StatusOK && ciLogKeyPattern.Match(body) {
			out <- recon.Finding{
				ProviderName: subdomain,
				Source:       probeURL,
				SourceType:   "recon:securitytrails",
				Confidence:   "high",
				DetectedAt:   time.Now(),
			}
		}
	}
}