package arxiv

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const atomNamespace = "http://www.w3.org/2005/Atom"

var titleWhitespace = regexp.MustCompile(`[\s\p{Z}]+`)

type atomFeed struct {
	XMLName      xml.Name    `xml:"feed"`
	TotalResults string      `xml:"totalResults"`
	StartIndex   string      `xml:"startIndex"`
	ItemsPerPage string      `xml:"itemsPerPage"`
	Entries      []atomEntry `xml:"entry"`
}

type atomEntry struct {
	ID              string         `xml:"id"`
	Updated         string         `xml:"updated"`
	Published       string         `xml:"published"`
	Title           string         `xml:"title"`
	Summary         string         `xml:"summary"`
	Authors         []atomAuthor   `xml:"author"`
	Comment         string         `xml:"comment"`
	JournalRef      string         `xml:"journal_ref"`
	DOI             string         `xml:"doi"`
	PrimaryCategory atomCategory   `xml:"primary_category"`
	Categories      []atomCategory `xml:"category"`
	Links           []atomLink     `xml:"link"`
}

type atomAuthor struct {
	Name         string   `xml:"name"`
	Affiliations []string `xml:"affiliation"`
}

type atomCategory struct {
	Term string `xml:"term,attr"`
}

type atomLink struct {
	Href  string `xml:"href,attr"`
	Title string `xml:"title,attr"`
	Rel   string `xml:"rel,attr"`
	Type  string `xml:"type,attr"`
}

func parseFeed(blob []byte) (SearchResult, error) {
	decoder := xml.NewDecoder(bytes.NewReader(blob))
	decoder.Strict = true
	var feed atomFeed
	if err := decoder.Decode(&feed); err != nil {
		return SearchResult{}, fmt.Errorf("decode Atom feed: %w", err)
	}
	if feed.XMLName.Local != "feed" || feed.XMLName.Space != atomNamespace {
		return SearchResult{}, fmt.Errorf("invalid Atom feed root")
	}
	total, err := parseFeedCount("totalResults", feed.TotalResults)
	if err != nil {
		return SearchResult{}, err
	}
	start, err := parseFeedCount("startIndex", feed.StartIndex)
	if err != nil {
		return SearchResult{}, err
	}
	if _, err := parseFeedCount("itemsPerPage", feed.ItemsPerPage); err != nil {
		return SearchResult{}, err
	}

	papers := make([]Paper, 0, len(feed.Entries))
	for _, entry := range feed.Entries {
		paper, ok := parseEntry(entry)
		if ok {
			papers = append(papers, paper)
		}
	}
	return SearchResult{
		Papers:       papers,
		TotalResults: total,
		Offset:       start,
		Returned:     len(papers),
		consumed:     len(feed.Entries),
	}, nil
}

func parseFeedCount(name, raw string) (int, error) {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return 0, fmt.Errorf("Atom feed missing %s", name)
	}
	value, err := strconv.Atoi(raw)
	if err != nil || value < 0 {
		return 0, fmt.Errorf("Atom feed has invalid %s", name)
	}
	return value, nil
}

func parseEntry(entry atomEntry) (Paper, bool) {
	entryURL, id, ok := normalizeArxivEntryURL(entry.ID)
	if !ok {
		return Paper{}, false
	}
	updated, err := time.Parse(time.RFC3339, strings.TrimSpace(entry.Updated))
	if err != nil {
		return Paper{}, false
	}
	published, err := time.Parse(time.RFC3339, strings.TrimSpace(entry.Published))
	if err != nil {
		return Paper{}, false
	}
	authors := make([]Author, 0, len(entry.Authors))
	for _, raw := range entry.Authors {
		affiliations := make([]string, 0, len(raw.Affiliations))
		for _, affiliation := range raw.Affiliations {
			affiliations = append(affiliations, affiliation)
		}
		authors = append(authors, Author{
			Name:         raw.Name,
			Affiliations: affiliations,
		})
	}
	links := make([]Link, 0, len(entry.Links))
	pdfURL := ""
	for _, raw := range entry.Links {
		href, ok := normalizedAbsoluteURL(raw.Href)
		if !ok {
			continue
		}
		links = append(links, Link{
			Href:        href,
			Title:       raw.Title,
			Rel:         raw.Rel,
			ContentType: raw.Type,
		})
		if pdfURL == "" && raw.Title == "pdf" {
			if candidate, valid := normalizeTrustedArxivArtifactURL(
				href,
				"/pdf/",
				id,
			); valid {
				pdfURL = candidate
			}
		}
	}
	categories := make([]string, 0, len(entry.Categories))
	for _, category := range entry.Categories {
		if category.Term != "" {
			categories = append(categories, category.Term)
		}
	}
	sourceURL := ""
	if pdfURL != "" {
		sourceURL = strings.Replace(pdfURL, "/pdf/", "/src/", 1)
	}
	return Paper{
		ID:              id,
		EntryURL:        entryURL,
		PublishedAt:     published.UTC(),
		UpdatedAt:       updated.UTC(),
		Title:           titleWhitespace.ReplaceAllString(entry.Title, " "),
		Authors:         authors,
		Abstract:        entry.Summary,
		Comment:         entry.Comment,
		JournalRef:      entry.JournalRef,
		DOI:             entry.DOI,
		PrimaryCategory: entry.PrimaryCategory.Term,
		Categories:      categories,
		Links:           links,
		PDFURL:          pdfURL,
		SourceURL:       sourceURL,
	}, true
}

func normalizeArxivEntryURL(raw string) (string, string, bool) {
	parsed, err := url.Parse(strings.TrimSpace(raw))
	if err != nil ||
		(parsed.Scheme != "http" && parsed.Scheme != "https") ||
		!trustedArxivHost(parsed) ||
		parsed.RawQuery != "" ||
		parsed.Fragment != "" {
		return "", "", false
	}
	const prefix = "/abs/"
	if !strings.HasPrefix(parsed.EscapedPath(), prefix) {
		return "", "", false
	}
	id, err := url.PathUnescape(
		strings.TrimPrefix(parsed.EscapedPath(), prefix),
	)
	if err != nil || !validPaperID(id) {
		return "", "", false
	}
	return "https://arxiv.org/abs/" + id, id, true
}

func normalizedAbsoluteURL(raw string) (string, bool) {
	parsed, err := url.Parse(strings.TrimSpace(raw))
	if err != nil ||
		(parsed.Scheme != "http" && parsed.Scheme != "https") ||
		parsed.Hostname() == "" ||
		parsed.User != nil {
		return "", false
	}
	if trustedArxivHost(parsed) {
		parsed.Scheme = "https"
	}
	return parsed.String(), true
}

func normalizeTrustedArxivArtifactURL(
	raw string,
	pathPrefix string,
	paperID string,
) (string, bool) {
	parsed, err := url.Parse(raw)
	if err != nil ||
		!trustedArxivHost(parsed) ||
		!strings.HasPrefix(parsed.EscapedPath(), pathPrefix) {
		return "", false
	}
	artifactID, err := url.PathUnescape(
		strings.TrimPrefix(parsed.EscapedPath(), pathPrefix),
	)
	if err != nil {
		return "", false
	}
	artifactID = strings.TrimSuffix(artifactID, ".pdf")
	if artifactID != paperID {
		return "", false
	}
	parsed.Scheme = "https"
	parsed.RawQuery = ""
	parsed.Fragment = ""
	return parsed.String(), true
}

func trustedArxivHost(candidate *url.URL) bool {
	if candidate == nil || candidate.User != nil || candidate.Port() != "" {
		return false
	}
	switch strings.ToLower(candidate.Hostname()) {
	case "arxiv.org", "www.arxiv.org", "export.arxiv.org":
		return true
	default:
		return false
	}
}