package criminalip

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"

	"errors"

	"github.com/projectdiscovery/uncover/sources"
)

const (
	URL        = "https://api.criminalip.io/v1/banner/search?query=%s&offset=%d"
	offsetStep = 10
	maxOffset  = 9900
)

type Agent struct{}

func (agent *Agent) Name() string {
	return "criminalip"
}

func (agent *Agent) Query(ctx context.Context, session *sources.Session, query *sources.Query) (chan sources.Result, error) {
	if session.Keys.CriminalIPToken == "" {
		return nil, errors.New("empty criminalip keys")
	}
	results := make(chan sources.Result)

	go func() {
		defer close(results)

		numberOfResults := 0
		currentPage := 0

		for {
			if ctx.Err() != nil {
				return
			}
			criminalipRequest := &CriminalIPRequest{
				Query:  query.Query,
				Offset: currentPage,
			}

			criminalipResponse := agent.query(ctx, URL, session, criminalipRequest, results)
			if criminalipResponse == nil {
				break
			}

			numberOfResults += len(criminalipResponse.Data.Result)

			nextOffset := currentPage + offsetStep

			if nextOffset > maxOffset {
				break
			}

			currentPage = nextOffset

			if numberOfResults > query.Limit || criminalipResponse.Data.Count == 0 || len(criminalipResponse.Data.Result) == 0 {
				break
			}
		}
	}()

	return results, nil
}

func (agent *Agent) queryURL(ctx context.Context, session *sources.Session, URL string, criminalipRequest *CriminalIPRequest) (*http.Response, error) {
	criminalipURL := fmt.Sprintf(URL, url.QueryEscape(criminalipRequest.Query), criminalipRequest.Offset)

	request, err := sources.NewHTTPRequest(ctx, http.MethodGet, criminalipURL, nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("x-api-key", session.Keys.CriminalIPToken)
	return session.Do(request, agent.Name())
}

func (agent *Agent) query(ctx context.Context, URL string, session *sources.Session, criminalipRequest *CriminalIPRequest, results chan sources.Result) *CriminalIPResponse {
	resp, err := agent.queryURL(ctx, session, URL, criminalipRequest)
	if err != nil {
		sources.SendResult(ctx, results, sources.Result{Source: agent.Name(), Error: err})
		return nil
	}

	criminalipResponse := &CriminalIPResponse{}
	if err := json.NewDecoder(resp.Body).Decode(criminalipResponse); err != nil {
		sources.SendResult(ctx, results, sources.Result{Source: agent.Name(), Error: err})
		return nil
	}
	if criminalipResponse.Status == http.StatusOK && criminalipResponse.Data.Count > 0 {
		for _, criminalipResult := range criminalipResponse.Data.Result {
			result := sources.Result{Source: agent.Name()}
			result.IP = criminalipResult.IP
			result.Port = criminalipResult.Port
			result.Host = criminalipResult.Domain
			raw, _ := json.Marshal(criminalipResult)
			result.Raw = raw
			if !sources.SendResult(ctx, results, result) {
				return criminalipResponse
			}
		}
	}

	return criminalipResponse
}

type CriminalIPRequest struct {
	Query  string
	Offset int
}