package main

import (
	"bytes"
	"database/sql"
	"fmt"
	"strings"
	"unicode"

	// "html/template"
	"net/http"
	"text/template"

	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/net/html"
)

type User struct {
	ID    int
	Name  string
	Posts []string
	Heros []string
}

var db *sql.DB

func initDB() {
	var err error
	db, err = sql.Open("sqlite3", "wormspace.db")
	if err != nil {
		panic(err)
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS users (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		name TEXT UNIQUE
	);`)
	if err != nil {
		panic(err)
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS posts (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		user_id INTEGER,
		content TEXT,
		FOREIGN KEY(user_id) REFERENCES users(id)
	);`)
	if err != nil {
		panic(err)
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS heros (
		user_id INTEGER,
		hero_id INTEGER,
		PRIMARY KEY(user_id, hero_id),
		FOREIGN KEY(user_id) REFERENCES users(id),
		FOREIGN KEY(hero_id) REFERENCES users(id)
	);`)
	if err != nil {
		panic(err)
	}
}

func setUser(w http.ResponseWriter, r *http.Request) {
	username := r.FormValue("username")
	if username == "" {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	for _, c := range username {
		if !unicode.IsLetter(c) {
			http.Error(w, "Invalid username", http.StatusBadRequest)
			return
		}
	}

	_, err := db.Exec("INSERT OR IGNORE INTO users (name) VALUES (?)", username)
	if err != nil {
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}

	http.SetCookie(w, &http.Cookie{Name: "username", Value: username, Path: "/"})
	http.Redirect(w, r, "/home", http.StatusSeeOther)
}

func getUser(r *http.Request) *User {
	cookie, err := r.Cookie("username")
	if err != nil {
		return nil
	}
	return &User{Name: cookie.Value}
}

func getPosts(username string) []string {
	rows, _ := db.Query("SELECT content FROM posts WHERE user_id = (SELECT id FROM users WHERE name = ?)", username)
	var posts []string
	for rows.Next() {
		var content string
		rows.Scan(&content)
		posts = append(posts, content)
	}
	return posts
}

func getHeros(username string) []string {
	rows, _ := db.Query("SELECT name FROM users WHERE id IN (SELECT hero_id FROM heros WHERE user_id = (SELECT id FROM users WHERE name = ?))", username)
	var heros []string
	for rows.Next() {
		var name string
		rows.Scan(&name)
		heros = append(heros, name)
	}
	return heros
}

// sanitizeHTML filters the input, allowing only <img> and <a> tags with "style" and "onload" attributes.
func sanitizeHTML(input string) string {
	doc, err := html.Parse(strings.NewReader(input))
	if err != nil {
		return ""
	}

	var buf bytes.Buffer
	processNode(&buf, doc)
	str := buf.String()
	cleaned := strings.ReplaceAll(str, "onreadystatechange", "")

	return cleaned
}

// processNode recursively processes nodes, allowing only specific elements and attributes.
func processNode(buf *bytes.Buffer, n *html.Node) {
	if n.Type == html.ElementNode {
		if n.Data != "img" && n.Data != "a" {
			// Skip non-allowed tags but still process children
			for c := n.FirstChild; c != nil; c = c.NextSibling {
				processNode(buf, c)
			}
			return
		}

		// Start tag
		buf.WriteString("<" + n.Data)

		// Filter attributes
		for _, attr := range n.Attr {
			if attr.Key == "onerror" || attr.Key == "src" || attr.Key == "href" {
				buf.WriteString(fmt.Sprintf(` %s="%s"`, attr.Key, attr.Val))
			}
		}

		buf.WriteString(">")

		// Process child nodes (for <a> which can have text)
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			processNode(buf, c)
		}

		// Close tag
		buf.WriteString("</" + n.Data + ">")
	} else if n.Type == html.TextNode {
		// Preserve text inside <a> tags
		buf.WriteString(n.Data)
	}

	// Process other children
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		processNode(buf, c)
	}
}

func addPost(w http.ResponseWriter, r *http.Request) {
	user := getUser(r)
	if user == nil {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	content := r.FormValue("post")
	content = sanitizeHTML(content)
	if content == "" {
		http.Redirect(w, r, "/home", http.StatusSeeOther)
		return
	}

	_, err := db.Exec("INSERT INTO posts (user_id, content) VALUES ((SELECT id FROM users WHERE name = ?), ?)", user.Name, content)
	if err != nil {
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}

	http.Redirect(w, r, "/home", http.StatusSeeOther)
}

func addHero(w http.ResponseWriter, r *http.Request) {
	user := getUser(r)
	if user == nil {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	hero := r.URL.Query().Get("hero")
	if hero == "" {
		http.Redirect(w, r, "/home", http.StatusSeeOther)
		return
	}

	_, err := db.Exec("INSERT OR IGNORE INTO heros (user_id, hero_id) VALUES ((SELECT id FROM users WHERE name = ?), (SELECT id FROM users WHERE name = ?))", user.Name, hero)
	if err != nil {
		http.Error(w, "Database error", http.StatusInternalServerError)
		return
	}

	http.Redirect(w, r, "/profile?name="+hero, http.StatusSeeOther)
}

func profilePage(w http.ResponseWriter, r *http.Request) {
	name := r.URL.Query().Get("name")
	if name == "" {
		http.NotFound(w, r)
		return
	}

	user := getUser(r)
	posts := getPosts(name)
	heros := getHeros(name)

	tmpl := template.Must(template.New("profile").Parse(`
		<html><body>
		<nav>
			<a href="/home">Home</a> | <a href="/profile?name={{.User.Name}}">My Profile</a>
		</nav>
		<h1>{{.Name}}'s Profile</h1>
		<a href="/addhero?hero={{.Name}}">Add as Hero</a>
		<h2>Posts:</h2>
		<ul>{{range .Posts}}<li>{{.}}</li>{{end}}</ul>
		<h2>Heros:</h2>
		<ul>{{range .Heros}}<li>{{.}}</li>{{end}}</ul>
		</body></html>`))
	tmpl.Execute(w, map[string]interface{}{"User": user, "Name": name, "Posts": posts, "Heros": heros})
}

func homePage(w http.ResponseWriter, r *http.Request) {
	user := getUser(r)
	if user == nil {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	rows, _ := db.Query("SELECT name FROM users")
	var users []string
	for rows.Next() {
		var name string
		rows.Scan(&name)
		users = append(users, name)
	}

	tmpl := template.Must(template.New("home").Parse(`
		<html><body>
		<nav>
			<a href="/home">Home</a> | <a href="/profile?name={{.User.Name}}">My Profile</a>
		</nav>
		<h1>Welcome {{.User.Name}}</h1>
		<form action="/addpost" method="POST">
			<input name="post" placeholder="Write a post">
			<input type="submit" value="Post">
		</form>
		<h2>Users:</h2>
		<ul>{{range .Users}}<li><a href="/profile?name={{.}}">{{.}}</a></li>{{end}}</ul>
		</body></html>`))
	tmpl.Execute(w, map[string]interface{}{"User": user, "Users": users})
}

func indexPage(w http.ResponseWriter, r *http.Request) {
	tmpl := template.Must(template.New("index").Parse(`
		<html><body>
		<h1>Welcome to Wormspace</h1>
		<form action="/setuser" method="POST">
			<input name="username" placeholder="Enter your username">
			<input type="submit" value="Login">
		</form>
		</body></html>`))
	tmpl.Execute(w, nil)
}

func main() {
	initDB()
	http.HandleFunc("/", indexPage)
	http.HandleFunc("/setuser", setUser)
	http.HandleFunc("/home", homePage)
	http.HandleFunc("/profile", profilePage)
	http.HandleFunc("/addpost", addPost)
	http.HandleFunc("/addhero", addHero)

	http.ListenAndServe(":8080", nil)
}