package httpx

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"strings"

	stringsutil "github.com/projectdiscovery/utils/strings"
	"golang.org/x/net/html"
	"slices"
)

var (
	cutset                  = "\n\t\v\f\r"
	reTitle                 = regexp.MustCompile(`(?im)<\s*title.*>(.*?)<\s*/\s*title>`)
	reContentType           = regexp.MustCompile(`(?im)\s*charset="(.*?)"|charset=(.*?)"\s*`)
	supportedTitleMimeTypes = []string{
		"text/html",
		"application/xhtml+xml",
		"application/xml",
		"application/rss+xml",
		"application/atom+xml",
		"application/xhtml+xml",
		"application/vnd.wap.xhtml+xml",
	}
)

// ExtractTitle from a response
func ExtractTitle(r *Response) (title string) {
	// Try to parse the DOM
	titleDom, err := getTitleWithDom(r)
	// In case of error fallback to regex
	if err != nil {
		for _, match := range reTitle.FindAllString(r.Raw, -1) {
			title = match
			break
		}
	} else {
		title = renderNode(titleDom)
	}

	title = html.UnescapeString(trimTitleTags(title))

	// remove unwanted chars
	title = strings.TrimSpace(strings.Trim(title, cutset))
	title = stringsutil.ReplaceAll(title, "", "\n", "\t", "\v", "\f", "\r")

	return title
}

func CanHaveTitleTag(mimeType string) bool {  
    return slices.Contains(supportedTitleMimeTypes, mimeType)  
}  

func getTitleWithDom(r *Response) (*html.Node, error) {
	var title *html.Node
	var crawler func(*html.Node)
	crawler = func(node *html.Node) {
		if node.Type == html.ElementNode && node.Data == "title" {
			title = node
			return
		}
		for child := node.FirstChild; child != nil && title == nil; child = child.NextSibling {
			crawler(child)
		}
	}
	htmlDoc, err := html.Parse(bytes.NewReader(r.Data))
	if err != nil {
		return nil, err
	}
	crawler(htmlDoc)
	if title != nil {
		return title, nil
	}
	return nil, fmt.Errorf("title not found")
}

func renderNode(n *html.Node) string {
	var buf bytes.Buffer
	w := io.Writer(&buf)
	html.Render(w, n) //nolint
	return buf.String()
}

func trimTitleTags(title string) string {
    titleBegin := strings.Index(title, ">")
    titleEnd := strings.Index(title, "</")
    if titleEnd < 0 || titleBegin < 0 || titleEnd <= titleBegin {
        return title
    }
    return title[titleBegin+1 : titleEnd]
}