package smtp

import (
	"context"
	"crypto/tls"
	"fmt"
	"github.com/pkg/errors"
	"net"
	"net/smtp"
	"net/url"
	"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)")
)

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

// Config contains the experiment config.
type Config struct{}

type RuntimeConfig struct {
	host       string
	port       string
	forced_tls bool
	noop_count uint8
}

func config(input model.MeasurementTarget) (*RuntimeConfig, 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 {
		// Valid port is checked by URL parsing
		port = parsed.Port()
	}

	valid_config := RuntimeConfig{
		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"`
	Runs    map[string]*IndividualTestKeys   `json:"runs"`
	// Used for global failure (DNS resolution)
	Failure string `json:"failure"`
	// Indicates global failure or individual test failure
	Failed bool `json:"failed"`
}

// IndividualTestKeys contains results for TCP/IP level stuff for each address found
// in the DNS lookup
type IndividualTestKeys struct {
	NoOpCounter   uint8
	TCPConnect    []*model.ArchivalTCPConnectResult         `json:"tcp_connect"`
	TLSHandshakes []*model.ArchivalTLSOrQUICHandshakeResult `json:"tls_handshakes"`
	// Individual failure aborting the test run for this address/port combo
	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 TCP sessions to the same hostname (over different IPs)
// don't use in parallel!
type TCPRunner struct {
	trace     *measurexlite.Trace
	logger    model.Logger
	ctx       context.Context
	tk        *TestKeys
	tlsconfig *tls.Config
	host      string
	port      string
	// addr is changed everytime TCPRunner.conn(addr) is called
	addr string
}

type TCPSession struct {
	addr     string
	port     string
	runner   *TCPRunner
	tk       *IndividualTestKeys
	tls      bool
	raw_conn *net.Conn
	tls_conn *net.Conn
}

func (s *TCPSession) Close() {
	if s.tls {
		var conn = *s.tls_conn
		conn.Close()
	} else {
		var conn = *s.raw_conn
		conn.Close()
	}
}

func (s *TCPSession) current_conn() net.Conn {
	if s.tls {
		return *s.tls_conn
	} else {
		return *s.raw_conn
	}
}

func (r *TCPRunner) run_key() string {
	return net.JoinHostPort(r.addr, r.port)
}

func (r *TCPRunner) get_run() *IndividualTestKeys {
	if r.tk.Runs == nil {
		r.tk.Runs = make(map[string]*IndividualTestKeys)
	}
	key := r.run_key()
	val, exists := r.tk.Runs[key]
	if exists {
		return val
	} else {
		r.tk.Runs[key] = &IndividualTestKeys{}
		return r.tk.Runs[key]
	}
}

func (r *TCPRunner) conn(addr string, port string) (*TCPSession, bool) {
	r.addr = addr
	run := r.get_run()

	s := new(TCPSession)
	if !s.conn(addr, port, r, run) {
		return nil, false
	}
	return s, true
}

func (r *TCPRunner) dial(addr string, port string) (net.Conn, error) {
	dialer := r.trace.NewDialerWithoutResolver(r.logger)
	conn, err := dialer.DialContext(r.ctx, "tcp", net.JoinHostPort(addr, port))
	run := r.get_run()
	run.TCPConnect = append(run.TCPConnect, r.trace.TCPConnects()...)
	return conn, err

}

func (s *TCPSession) conn(addr string, port string, runner *TCPRunner, tk *IndividualTestKeys) bool {
	// Initialize addr field and corresponding errors in TestKeys
	s.addr = addr
	s.port = port
	s.tls = false
	s.runner = runner
	s.tk = tk

	conn, err := runner.dial(addr, port)
	if err != nil {
		s.error(err)
		return false
	}
	s.raw_conn = &conn

	return true
}

func (s *TCPSession) error(err error) {
	s.runner.tk.Failed = true
	s.tk.Failure = tracex.NewFailure(err)
	//s. = append(s.errors, tracex.NewFailure(err))
}

func (r *TCPRunner) 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 (s *TCPSession) handshake() bool {
	if s.tls {
		// TLS already initialized...
		return true
	}
	s.runner.logger.Infof("Starting TLS handshake with %s:%s", s.addr, s.port)
	thx := s.runner.trace.NewTLSHandshakerStdlib(s.runner.logger)
	tconn, _, err := thx.Handshake(s.runner.ctx, *s.raw_conn, s.runner.tlsconfig)
	s.tk.TLSHandshakes = append(s.tk.TLSHandshakes, s.runner.trace.FirstTLSHandshakeOrNil())
	if err != nil {
		s.error(err)
		return false
	}

	s.tls = true
	s.tls_conn = &tconn
	s.runner.logger.Infof("Handshake succeeded")
	return true
}

func (s *TCPSession) starttls(message string) bool {
	if s.tls {
		// TLS already initialized...
		return true
	}
	if message != "" {
		s.runner.logger.Infof("Asking for StartTLS upgrade")
		s.current_conn().Write([]byte(message))
	}
	return s.handshake()
}

func (s *TCPSession) smtp(ehlo string, noop uint8) bool {
	// Auto-choose plaintext/TCP session
	client, err := smtp.NewClient(s.current_conn(), ehlo)
	if err != nil {
		s.error(err)
		return false
	}
	err = client.Hello(ehlo)
	if err != nil {
		s.error(err)
		return false
	}

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

		if s.tk.NoOpCounter == noop {
			s.runner.logger.Infof("Successfully generated no-op traffic")
			return true
		} else {
			s.runner.logger.Infof("Failed no-op traffic at iteration %d", s.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)
	measurement.TestKeys = tk

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

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

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

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

	for _, addr := range addrs {
		tcp_session, success := runner.conn(addr, config.port)
		if !success {
			continue
		}
		defer tcp_session.Close()

		if config.forced_tls {
			// Direct TLS connection
			if !tcp_session.handshake() {
				continue
			}

			// Try EHLO + NoOps
			if !tcp_session.smtp("localhost", config.noop_count) {
				continue
			}
		} else {
			// StartTLS... first try plaintext EHLO
			if !tcp_session.smtp("localhost", 0) {
				continue
			}

			// Upgrade via StartTLS and try EHLO + NoOps
			if !tcp_session.starttls("STARTTLS\n") {
				continue
			}

			if !tcp_session.smtp("localhost", config.noop_count) {
				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
}