// Package discovery implements the server discovery by contacting disco.eduvpn.org and returning the data as a Go structure
package discovery

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"codeberg.org/eduVPN/eduvpn-common/internal/httpwrap"
	"codeberg.org/eduVPN/eduvpn-common/internal/levenshtein"
	"codeberg.org/eduVPN/eduvpn-common/internal/verify"
	discotypes "codeberg.org/eduVPN/eduvpn-common/types/discovery"

	_ "embed"
)

// eServers are the embedded server list
//
//go:embed server_list.json
var eServers []byte

// eOrganizations are the embedded organizations
//
//go:embed organization_list.json
var eOrganizations []byte

// Organizations are the list of organizations from https://disco.eduvpn.org/v2/organization_list.json
type Organizations struct {
	// Version is the version field in discovery. The Go library checks this for rollbacks
	Version uint64 `json:"v"`
	// List is the list of organizations, omitted if empty
	List []Organization `json:"organization_list,omitempty"`
	// Timestamp is the timestamp that is internally used by the Go library to keep track
	// of when the organizations were last updated
	Timestamp time.Time `json:"go_timestamp"`
	// UpdateHeader is the result of the "Last-Modified" header
	UpdateHeader time.Time `json:"go_update_header"`
}

// Organization is a single discovery Organization
type Organization struct {
	// Organization is the embedded public type that is a subset of this thus common Organization
	discotypes.Organization
	// SecureInternetHome is the secure internet home server that belongs to this organization
	// Omitted if none is defined
	SecureInternetHome string `json:"secure_internet_home"`
	// KeywordList is the list of keywords
	// Omitted if none is defined
	KeywordList discotypes.MapOrString `json:"keyword_list,omitempty"`
}

// Score returns the levenshstein score for a search query `search` on the organizations
func (o *Organization) Score(search string) int {
	return levenshtein.DiscoveryScore(search, o.DisplayName, o.KeywordList)
}

// Servers are the list of servers from https://disco.eduvpn.org/v2/server_list.json
type Servers struct {
	// Version is the version field in discovery. The Go library checks this for rollbacks
	Version uint64 `json:"v"`
	// List is the list of servers, omitted if empty
	List []Server `json:"server_list,omitempty"`
	// Timestamp is a timestamp that is internally used by the Go library to keek track
	// of when the servers were last updated
	Timestamp time.Time `json:"go_timestamp"`
	// UpdateHeader is the result of the "Last-Modified" header
	UpdateHeader time.Time `json:"go_update_header"`
}

// Server is a single discovery server
type Server struct {
	// Server is the embedded public type that is a subset of this common Server
	discotypes.Server
	// AuthenticationURLTemplate is the template to be used for authentication to skip WAYF
	AuthenticationURLTemplate string `json:"authentication_url_template,omitempty"`
	// KeywordList are the keywords of the server, omitted if empty
	KeywordList discotypes.MapOrString `json:"keyword_list,omitempty"`
	// PublicKeyList are the public keys of the server. Currently not used in this lib but returned by the upstream discovery server
	PublicKeyList []string `json:"public_key_list,omitempty"`
	// SupportContact is the list/slice of support contacts
	SupportContact []string `json:"support_contact,omitempty"`
}

// Score returns the score of the search query `str` on this server
func (s *Server) Score(search string) int {
	return levenshtein.DiscoveryScore(search, s.DisplayName, s.KeywordList)
}

// Discovery is the main structure used for this package.
type Discovery struct {
	// mu is the read write mutex that protects the struct from concurrent access
	mu sync.RWMutex
	// The httpClient for sending HTTP requests
	httpClient *httpwrap.Client

	// Organizations represents the organizations that are returned by the discovery server
	OrganizationList Organizations `json:"organizations"`

	// Servers represents the servers that are returned by the discovery server
	ServerList Servers `json:"servers"`
}

// DiscoURL is the URL used for fetching the discovery files and signatures
var DiscoURL = "https://disco.eduvpn.org/v2/"

// file is a helper function that gets a disco JSON and fills the structure with it
// If it was unsuccessful it returns an error.
func (discovery *Discovery) file(ctx context.Context, jsonFile string, previousVersion uint64, last time.Time, structure any) (time.Time, error) {
	var newUpdate time.Time
	// No HTTP client present, create one
	if discovery.httpClient == nil {
		discovery.httpClient = httpwrap.NewClient(nil)
	}

	// Get json data
	jsonURL, err := httpwrap.JoinURLPath(DiscoURL, jsonFile)
	if err != nil {
		return newUpdate, err
	}

	var opts *httpwrap.OptionalParams
	if !last.IsZero() {
		header := http.Header{
			"If-Modified-Since": []string{last.Format(http.TimeFormat)},
		}
		opts = &httpwrap.OptionalParams{
			Headers: header,
		}
	}
	h, body, err := discovery.httpClient.Do(ctx, "GET", jsonURL, opts)
	if err != nil {
		return newUpdate, err
	}

	lms := h.Get("Last-Modified")
	if lms != "" {
		lm, err := http.ParseTime(lms)
		if err != nil {
			slog.Warn("failed to parse 'Last-Modified' header", "error", err)
		} else {
			newUpdate = lm
			slog.Debug("got 'Last-Modified' header", "value", lm)
		}
	} else {
		slog.Warn("no 'Last-Modified' header found")
	}

	// Get signature
	sigFile := jsonFile + ".minisig"
	sigURL, err := httpwrap.JoinURLPath(DiscoURL, sigFile)
	if err != nil {
		return newUpdate, err
	}
	_, sigBody, err := discovery.httpClient.Get(ctx, sigURL)
	if err != nil {
		return newUpdate, err
	}

	// Verify signature
	// Set this to true when we want to force prehash
	const forcePrehash = false
	ok, err := verify.Verify(
		string(sigBody),
		body,
		jsonFile,
		previousVersion,
		forcePrehash,
	)

	if !ok || err != nil {
		return newUpdate, err
	}

	// Parse JSON to extract version and list
	if err = json.Unmarshal(body, structure); err != nil {
		return newUpdate, fmt.Errorf("failed parsing discovery file: '%s' from the server with error: %w and body: %s", jsonFile, err, string(body))
	}

	return newUpdate, nil
}

// MarkOrganizationsExpired marks the organizations as expired
func (discovery *Discovery) MarkOrganizationsExpired() {
	discovery.mu.Lock()
	defer discovery.mu.Unlock()
	// Re-initialize the timestamp to zero
	discovery.OrganizationList.Timestamp = time.Time{}
}

// MarkServersExpired marks the servers as expired
func (discovery *Discovery) MarkServersExpired() {
	discovery.mu.Lock()
	defer discovery.mu.Unlock()
	// Re-initialize the timestamp to zero
	discovery.ServerList.Timestamp = time.Time{}
}

// determineOrganizationsUpdate returns a boolean indicating whether or not the discovery organizations should be updated
// https://github.com/eduvpn/documentation/blob/v3/SERVER_DISCOVERY.md
// - [IMPLEMENTED] on "first launch" when offering the search for "Institute Access" and "Organizations";
// - [IMPLEMENTED in client/client.go and here] when the user tries to add new server AND the user did NOT yet choose an organization before; Implemented in Register()
// - [IMPLEMENTED in client/client.go] when the authorization for the server associated with an already chosen organization is triggered, e.g. after expiry or revocation.
// - [IMPLEMENTED here] NOTE: when the org_id that the user chose previously is no longer available in organization_list.json the application should ask the user to choose their organization (again). This can occur for example when the organization replaced their identity provider, uses a different domain after rebranding or simply ceased to exist.
func (discovery *Discovery) determineOrganizationsUpdate() bool {
	if discovery.OrganizationList.Timestamp.IsZero() {
		return true
	}
	if discovery.OrganizationList.UpdateHeader.IsZero() {
		return true
	}
	return false
}

// SecureLocationList returns a slice of all the available locations.
func (discovery *Discovery) SecureLocationList() []string {
	discovery.mu.RLock()
	defer discovery.mu.RUnlock()
	var loc []string
	for _, srv := range discovery.ServerList.List {
		if srv.Type == "secure_internet" {
			loc = append(loc, srv.CountryCode)
		}
	}
	return loc
}

// ServerByURL returns the discovery server by the base URL and the according type ("secure_internet", "institute_access")
// An error is returned if and only if nil is returned for the server.
func (discovery *Discovery) ServerByURL(
	baseURL string,
	srvType string,
) (*Server, error) {
	discovery.mu.RLock()
	defer discovery.mu.RUnlock()
	for _, currentServer := range discovery.ServerList.List {
		if currentServer.BaseURL == baseURL && currentServer.Type == srvType {
			return &currentServer, nil
		}
	}
	return nil, fmt.Errorf("no server of type '%s' at URL '%s'", srvType, baseURL)
}

// ErrCountryNotFound is used when the secure internet country cannot be found
type ErrCountryNotFound struct {
	CountryCode string
}

func (cnf *ErrCountryNotFound) Error() string {
	return fmt.Sprintf("no secure internet server with country code: '%s'", cnf.CountryCode)
}

// ServerByCountryCode returns the discovery server by the country code
// An error is returned if and only if nil is returned for the server.
func (discovery *Discovery) ServerByCountryCode(countryCode string) (*Server, error) {
	discovery.mu.RLock()
	defer discovery.mu.RUnlock()
	for _, srv := range discovery.ServerList.List {
		if srv.CountryCode == countryCode && srv.Type == "secure_internet" {
			return &srv, nil
		}
	}
	return nil, &ErrCountryNotFound{CountryCode: countryCode}
}

// orgByID returns the discovery organization by the organization ID
// An error is returned if and only if nil is returned for the organization.
func (discovery *Discovery) orgByID(orgID string) (*Organization, error) {
	for _, org := range discovery.OrganizationList.List {
		if org.OrgID == orgID {
			return &org, nil
		}
	}
	return nil, fmt.Errorf("no secure internet home found in organization '%s'", orgID)
}

// SecureHomeArgs returns the secure internet home server arguments:
// - The organization it belongs to
// - The secure internet server itself
// An error is returned if and only if nil is returned for the organization.
func (discovery *Discovery) SecureHomeArgs(orgID string) (*Organization, *Server, error) {
	discovery.mu.RLock()
	org, err := discovery.orgByID(orgID)
	if err != nil {
		discovery.mu.RUnlock()
		discovery.MarkOrganizationsExpired()
		return nil, nil, err
	}

	// Get a server with the base url
	srv, err := discovery.ServerByURL(org.SecureInternetHome, "secure_internet")
	if err != nil {
		discovery.mu.RUnlock()
		discovery.MarkOrganizationsExpired()
		return nil, nil, err
	}
	discovery.mu.RUnlock()
	return org, srv, nil
}

// determineServersUpdate returns whether or not the discovery servers should be updated by contacting the discovery server
// https://github.com/eduvpn/documentation/blob/v3/SERVER_DISCOVERY.md
// - [Implemented] The application MUST always fetch the server_list.json at application start.
// - The application MAY refresh the server_list.json periodically, e.g. once every hour.
func (discovery *Discovery) determineServersUpdate() bool {
	// No servers, we should update
	if discovery.ServerList.Timestamp.IsZero() {
		return true
	}
	if discovery.ServerList.UpdateHeader.IsZero() {
		return true
	}
	// 1 hour from the last update
	upd := discovery.ServerList.Timestamp.Add(1 * time.Hour)
	return !time.Now().Before(upd)
}

func (discovery *Discovery) cachedOrgs() *Organizations {
	if discovery.OrganizationList.Version == 0 {
		return nil
	}
	return &discovery.OrganizationList
}

// Organizations returns the discovery organizations
// The second return value is a boolean that indicates whether a fresh list was updated internally
// If there was an error, a cached copy is returned if available.
// cache is set to true if there should be no network call done
func (discovery *Discovery) Organizations(ctx context.Context, cache bool) (*Organizations, bool, error) {
	discovery.mu.RLock()
	if cache || !discovery.determineOrganizationsUpdate() {
		discovery.mu.RUnlock()
		return discovery.cachedOrgs(), false, nil
	}
	discovery.mu.RUnlock()
	discovery.mu.Lock()
	defer discovery.mu.Unlock()
	file := "organization_list.json"
	var jsonDecode Organizations
	update, err := discovery.file(ctx, file, discovery.OrganizationList.Version, discovery.OrganizationList.UpdateHeader, &jsonDecode)
	if err != nil {
		statErr := &httpwrap.StatusError{}
		if errors.As(err, &statErr) {
			if statErr.Status != 304 {
				slog.Warn("failed to get fresh organization", "error", err)
			} else {
				discovery.OrganizationList.Timestamp = time.Now()
				slog.Debug("got HTTP 304 for discovery, organization_list.json is not modified")
				err = nil
			}
		}
		// Return previous with an error
		return discovery.cachedOrgs(), false, err
	}
	if len(jsonDecode.List) == 0 {
		slog.Warn("fresh organization list is empty")
	} else {
		discovery.OrganizationList = jsonDecode
	}
	discovery.OrganizationList.Timestamp = time.Now()
	if !update.IsZero() {
		discovery.OrganizationList.UpdateHeader = update
	}
	return &discovery.OrganizationList, true, nil
}

func (discovery *Discovery) cachedServers() *Servers {
	if discovery.ServerList.Version == 0 {
		return nil
	}
	return &discovery.ServerList
}

// Servers returns the discovery servers
// The second return value is a boolean that indicates whether a fresh list was updated internally
// If there was an error, a cached copy is returned if available.
// cache is set to true if there should be no network call done
func (discovery *Discovery) Servers(ctx context.Context, cache bool) (*Servers, bool, error) {
	discovery.mu.RLock()
	if cache || !discovery.determineServersUpdate() {
		discovery.mu.RUnlock()
		return discovery.cachedServers(), false, nil
	}
	discovery.mu.RUnlock()
	discovery.mu.Lock()
	defer discovery.mu.Unlock()
	file := "server_list.json"
	var jsonDecode Servers
	update, err := discovery.file(ctx, file, discovery.ServerList.Version, discovery.ServerList.UpdateHeader, &jsonDecode)
	if err != nil {
		statErr := &httpwrap.StatusError{}
		if errors.As(err, &statErr) {
			if statErr.Status != 304 {
				slog.Warn("failed to get fresh servers", "error", err)
			} else {
				discovery.ServerList.Timestamp = time.Now()
				slog.Debug("got HTTP 304 for discovery, server_list.json is not modified")
				err = nil
			}
		}
		// Return previous with an error
		return discovery.cachedServers(), false, err
	}
	if len(jsonDecode.List) == 0 {
		slog.Warn("fresh server list is empty")
	} else {
		discovery.ServerList = jsonDecode
	}
	discovery.ServerList.Timestamp = time.Now()
	if !update.IsZero() {
		discovery.ServerList.UpdateHeader = update
	}
	return &discovery.ServerList, true, nil
}

// UpdateServers updates the discovery servers to the new version
// It does this by checking versions
func (discovery *Discovery) UpdateServers(other Servers) {
	if other.Version >= discovery.ServerList.Version {
		discovery.ServerList = other
	}
}

// UpdateOrganizations updates the discovery organizations to the new version
// It does this by checking versions
func (discovery *Discovery) UpdateOrganizations(other Organizations) {
	if other.Version >= discovery.OrganizationList.Version {
		discovery.OrganizationList = other
	}
}

// Fill makes sure that the cache is filled with the embedded discovery
func (discovery *Discovery) Fill() error {
	discovery.mu.Lock()
	defer discovery.mu.Unlock()

	var err error
	var es Servers
	err = errors.Join(err, json.Unmarshal(eServers, &es))
	discovery.UpdateServers(es)
	var eo Organizations
	err = errors.Join(err, json.Unmarshal(eOrganizations, &eo))
	discovery.UpdateOrganizations(eo)
	return err
}
