package smtp

import (
	"context"
	"crypto/tls"
	"fmt"
	"github.com/pkg/errors"
	"net"
	"net/smtp"
	"net/url"
	"strconv"
	"time"

	"github.com/ooni/probe-cli/v3/internal/engine/experiment/urlgetter"
	"github.com/ooni/probe-cli/v3/internal/measurexlite"
	"github.com/ooni/probe-cli/v3/internal/model"
	"github.com/ooni/probe-cli/v3/internal/tracex"
)

var (
	// errNoInputProvided indicates you didn't provide any input
	errNoInputProvided = errors.New("not input provided")

	// errInputIsNotAnURL indicates that input is not an URL
	errInputIsNotAnURL = errors.New("input is not an URL")

	// errInvalidScheme indicates that the scheme is invalid
	errInvalidScheme = errors.New("scheme must be smtp(s)")

	// errInvalidPort indicates that the port provided could not be parsed as an int
	errInvalidPort = errors.New("Port number is not a valid integer")
)

const (
	testName    = "smtp"
	testVersion = "0.0.1"
)

// Config contains the experiment config.
type Config struct {
	host       string
	port       string
	forced_tls bool
	noop_count uint8
}

func config(input model.MeasurementTarget) (*Config, error) {
	if input == "" {
		// TODO: static input data (eg. gmail/riseup..)
		return nil, errNoInputProvided
	}

	parsed, err := url.Parse(string(input))
	if err != nil {
		return nil, fmt.Errorf("%w: %s", errInputIsNotAnURL, err.Error())
	}
	if parsed.Scheme != "smtp" && parsed.Scheme != "smtps" {
		return nil, errInvalidScheme
	}

	port := ""

	if parsed.Port() == "" {
		// Default ports for StartTLS and forced TLS respectively
		if parsed.Scheme == "smtp" {
			port = "587"
		} else {
			port = "465"
		}
	} else {
		// Check that requested port is a valid integer
		_, err := strconv.Atoi(parsed.Port())
		if err != nil {
			return nil, errInvalidPort
		} else {
			port = parsed.Port()
		}
	}

	valid_config := Config{
		host:       parsed.Hostname(),
		forced_tls: parsed.Scheme == "smtps",
		port:       port,
		noop_count: 10,
	}

	return &valid_config, nil
}

// TestKeys contains the experiment results
type TestKeys struct {
	Queries       []*model.ArchivalDNSLookupResult          `json:"queries"`
	TCPConnect    []*model.ArchivalTCPConnectResult         `json:"tcp_connect"`
	TLSHandshakes []*model.ArchivalTLSOrQUICHandshakeResult `json:"tls_handshakes"`
	SMTPErrors    map[string][]*string                      `json:"smtp"`
	NoOpCounter   uint8                                     `json:"successful_noops"`
	// Used for global failure (DNS resolution)
	Failure string `json:"failure"`
}

type Measurer struct {
	// Config contains the experiment settings. If empty we
	// will be using default settings.
	Config Config

	// Getter is an optional getter to be used for testing.
	Getter urlgetter.MultiGetter
}

// ExperimentName implements ExperimentMeasurer.ExperimentName
func (m Measurer) ExperimentName() string {
	return testName
}

// ExperimentVersion implements ExperimentMeasurer.ExperimentVersion
func (m Measurer) ExperimentVersion() string {
	return testVersion
}

// Manages sequential SMTP sessions to the same hostname (over different IPs)
// don't use in parallel!
type SMTPRunner struct {
	trace     *measurexlite.Trace
	logger    model.Logger
	ctx       context.Context
	tk        *TestKeys
	tlsconfig *tls.Config
	host      string
	port      string
	// addr is changed everytime SMTPRunner.conn(addr) is called
	addr string
}

func (r SMTPRunner) smtp_error(err error) {
	key := net.JoinHostPort(r.addr, r.port)
	errors, exists := r.tk.SMTPErrors[key]
	if exists {
		r.tk.SMTPErrors[key] = append(errors, tracex.NewFailure(err))
	} else {
		r.tk.SMTPErrors[key] = []*string{tracex.NewFailure(err)}
	}
}

func (r SMTPRunner) resolve(host string) ([]string, bool) {
	r.logger.Infof("Resolving DNS for %s", host)
	resolver := r.trace.NewStdlibResolver(r.logger)
	addrs, err := resolver.LookupHost(r.ctx, host)
	r.tk.Queries = append(r.tk.Queries, r.trace.DNSLookupsFromRoundTrip()...)
	if err != nil {
		r.tk.Failure = *tracex.NewFailure(err)
		return []string{}, false
	}
	r.logger.Infof("Finished DNS for %s: %v", host, addrs)

	return addrs, true
}

func (r SMTPRunner) conn(addr string) (net.Conn, bool) {
	r.addr = addr
	dialer := r.trace.NewDialerWithoutResolver(r.logger)
	conn, err := dialer.DialContext(r.ctx, "tcp", net.JoinHostPort(r.addr, r.port))
	r.tk.TCPConnect = append(r.tk.TCPConnect, r.trace.TCPConnects()...)
	if err != nil {
		r.smtp_error(err)
		return nil, false
	}
	return conn, true
}

func (r SMTPRunner) handshake(conn net.Conn) (net.Conn, bool) {
	r.logger.Infof("Starting TLS handshake with %s:%s (%s)", r.host, r.port, r.addr)
	thx := r.trace.NewTLSHandshakerStdlib(r.logger)
	tconn, _, err := thx.Handshake(r.ctx, conn, r.tlsconfig)
	r.tk.TLSHandshakes = append(r.tk.TLSHandshakes, r.trace.FirstTLSHandshakeOrNil())
	if err != nil {
		r.smtp_error(err)
		return nil, false
	}
	r.logger.Infof("Handshake succeeded")
	return tconn, true
}

func (r SMTPRunner) starttls(conn net.Conn, message string) (net.Conn, bool) {
	if message != "" {
		r.logger.Infof("Asking for StartTLS upgrade")
		conn.Write([]byte(message))
	}
	tconn, success := r.handshake(conn)
	return tconn, success
}

func (r SMTPRunner) smtp(conn net.Conn, ehlo string, noop uint8) bool {
	client, err := smtp.NewClient(conn, ehlo)
	if err != nil {
		r.smtp_error(err)
		return false
	}
	err = client.Hello(ehlo)
	if err != nil {
		r.smtp_error(err)
		return false
	}

	if noop > 0 {
		r.logger.Infof("Trying to generate more no-op traffic")
		// TODO: noop counter per IP address
		r.tk.NoOpCounter = 0
		for r.tk.NoOpCounter < noop {
			r.tk.NoOpCounter += 1
			r.logger.Infof("NoOp Iteration %d", r.tk.NoOpCounter)
			err = client.Noop()
			if err != nil {
				r.smtp_error(err)
				break
			}
		}

		if r.tk.NoOpCounter == noop {
			r.logger.Infof("Successfully generated no-op traffic")
			return true
		} else {
			r.logger.Infof("Failed no-op traffic at iteration %d", r.tk.NoOpCounter)
			return false
		}
	}

	return true
}

// Run implements ExperimentMeasurer.Run
func (m Measurer) Run(
	ctx context.Context, sess model.ExperimentSession,
	measurement *model.Measurement, callbacks model.ExperimentCallbacks,
) error {

	log := sess.Logger()
	trace := measurexlite.NewTrace(0, measurement.MeasurementStartTimeSaved)

	config, err := config(measurement.Input)
	if err != nil {
		// Invalid input data, we don't even generate report
		return err
	}

	tk := new(TestKeys)
	// TODO: make it so we don't forget it
	tk.SMTPErrors = make(map[string][]*string)
	measurement.TestKeys = tk

	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	tlsconfig := tls.Config{
		InsecureSkipVerify: false,
		ServerName:         config.host,
	}

	runner := SMTPRunner{
		trace:     trace,
		logger:    log,
		ctx:       ctx,
		tk:        tk,
		tlsconfig: &tlsconfig,
		host:      config.host,
		port:      config.port,
		addr:      "",
	}

	// First resolve DNS
	addrs, success := runner.resolve(config.host)
	if !success {
		return nil
	}

	for _, addr := range addrs {
		conn, success := runner.conn(addr)
		if !success {
			return nil
		}
		defer conn.Close()

		if config.forced_tls {
			// Direct TLS connection
			tconn, success := runner.handshake(conn)
			if !success {
				continue
			}
			defer tconn.Close()

			// Try EHLO + NoOps
			if !runner.smtp(tconn, "localhost", 10) {
				continue
			}
		} else {
			// StartTLS... first try plaintext EHLO
			if !runner.smtp(conn, "localhost", 0) {
				continue
			}

			// Upgrade via StartTLS and try EHLO + NoOps
			tconn, success := runner.starttls(conn, "STARTTLS\n")
			if !success {
				continue
			}
			defer tconn.Close()

			if !runner.smtp(tconn, "localhost", 10) {
				continue
			}
		}
	}

	return nil
}

// NewExperimentMeasurer creates a new ExperimentMeasurer.
func NewExperimentMeasurer(config Config) model.ExperimentMeasurer {
	return Measurer{Config: config}
}

// SummaryKeys contains summary keys for this experiment.
//
// Note that this structure is part of the ABI contract with ooniprobe
// therefore we should be careful when changing it.
type SummaryKeys struct {
	//DNSBlocking bool `json:"facebook_dns_blocking"`
	//TCPBlocking bool `json:"facebook_tcp_blocking"`
	IsAnomaly bool `json:"-"`
}

// GetSummaryKeys implements model.ExperimentMeasurer.GetSummaryKeys.
func (m Measurer) GetSummaryKeys(measurement *model.Measurement) (interface{}, error) {
	sk := SummaryKeys{IsAnomaly: false}
	_, ok := measurement.TestKeys.(*TestKeys)
	if !ok {
		return sk, errors.New("invalid test keys type")
	}
	return sk, nil
}