server/main.go

package main

import (
	"database/sql"
	"encoding/xml"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"

	_ "github.com/mattn/go-sqlite3"
	"github.com/op/go-logging"

	toml "github.com/pelletier/go-toml/v2"
)

type Config struct {
	Address       string
	Port          int
	ScanDirectory string
	//Age        int
	//Cats       []string
	//Pi         float64
	//Perfection []int
	//DOB        time.Time // requires `import time`
}

func get_config_dir() string {
	// Find config directory

	// Linux (XDG base directory specification compliant)
	xdg_config_home := os.Getenv("XDG_CONFIG_HOME")

	if xdg_config_home != "" {
		return filepath.Join(xdg_config_home, "ikinuki")
	} else {
		home := os.Getenv("HOME")
		return filepath.Join(home, ".config", "ikinuki")
	}
}

func get_config(config_dir string) (*Config, error) {
	config := Config{
		Address:       "127.0.0.1",
		Port:          32520,
		ScanDirectory: "FILL IN",
	}

	err := os.MkdirAll(config_dir, 0755)
	if err != nil {
		return &config, fmt.Errorf("Could not create config directory: %v", err)
	}

	config_path := filepath.Join(config_dir, "server.toml")

	log.Debugf("Config File Dir: %v\n", config_path)

	// open file, creating it if empty
	config_file, err := os.OpenFile(config_path, os.O_RDWR|os.O_CREATE, 0755)
	defer config_file.Close()

	if err != nil {
		return &config, fmt.Errorf("Could not open config file: %v", err)
	}

	fi, err := config_file.Stat()
	if err != nil {
		return &config, fmt.Errorf("Could not get config file size: %v", err)
	}

	if fi.Size() == 0 {
		// if file is empty, write default config
		data, err := toml.Marshal(config)
		if err != nil {
			return &config, fmt.Errorf("Could not write default config to file: %v", err)
		}
		config_file.Write(data)
		// return nil as the user must edit the config file
		return nil, nil
	} else {
		// try to parse config file
		data, err := ioutil.ReadAll(config_file)
		if err != nil {
			return &config, fmt.Errorf("Could not read data from config file: %v", err)
		}
		err = toml.Unmarshal(data, &config)
		if err != nil {
			return &config, fmt.Errorf("Could not parse config file contents: %v", err)
		}
		return &config, nil
	}

}

func get_data_dir() string {
	// Find data directory

	// Linux (XDG base directory specification compliant)
	xdg_data_home := os.Getenv("XDG_DATA_HOME")

	if xdg_data_home != "" {
		return filepath.Join(xdg_data_home, "ikinuki")
	} else {
		home := os.Getenv("HOME")
		return filepath.Join(home, ".local", "share", "ikinuki")
	}
}

func open_database(data_dir string) (*sql.DB, error) {
	// create data directory
	err := os.MkdirAll(data_dir, 0755)
	if err != nil {
		return nil, fmt.Errorf("Could not create data directory: %v", err)
	}

	// open database
	db_path := filepath.Join(data_dir, "ikinuki-server.db")

	db, err := sql.Open("sqlite3", db_path)
	if err != nil {
		return nil, fmt.Errorf("Could not open database: %v", err)
	}

	// check if the database is new
	is_db_new := `SELECT CASE WHEN EXISTS(SELECT 1 FROM sqlite_master) THEN 0 ELSE 1 END AS IsEmpty;`

	// initalize database
	db_init := `
CREATE TABLE tvshows(
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    title TEXT,
    original_title TEXT,
    show_title TEXT,
    year INTEGER
);
CREATE TABLE episodes(
    id INTEGER PRIMARY KEY AUTOINCREMENT,
	title TEXT,
	show_title TEXT,
    number INTEGER,
    season INTEGER,
    original_filename TEXT
);
`

	// transaction new check isn't racy
	tx, err := db.Begin()
	if err != nil {
		return nil, fmt.Errorf("Could not begin db transaction: %v", err)
	}
	defer tx.Commit()

	// check if the database is new
	rows, err := tx.Query(is_db_new)
	if err != nil {
		return nil, fmt.Errorf("Could not check if database was new: %v", err)
	}
	defer rows.Close()

	rows.Next()
	var isEmpty int
	rows.Scan(&isEmpty)

	// if the sqlite_master table is empty (the databse is newly created)
	if isEmpty == 1 {
		log.Infof("Creating new database at %v\n", db_path)
		// initalize database
		_, err = tx.Exec(db_init)
		if err != nil {
			return nil, fmt.Errorf("Could not initalize database: %v", err)
		}
	} else {
		log.Infof("Existing database found at %v", db_path)
	}

	return db, nil
}

type TVShow struct {
	XMLName      xml.Name `xml:"tvshow" json:"-"`
	Title        string   `xml:"title"`
	OriginalTile string   `xml:"originaltitle"`
	ShowTitle    string   `xml:"showtitle"`
	Year         int      `xml:"year"`
}

func scan_tvshow_root(root string, db *sql.DB) error {
	log.Noticef("Scanning directory root %v", root)

	files, err := ioutil.ReadDir(root)
	if err != nil {
		return fmt.Errorf("Could not get files in directory \"%v\": %v", root, err)
	}

	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("Could not begin db transaction: %v", err)
	}
	defer tx.Commit()

	stmt, err := tx.Prepare("insert into tvshows(title, original_title, show_title, year) values(?, ?, ?, ?)")
	if err != nil {
		return fmt.Errorf("Could not prepare statement: %v", err)
	}
	defer stmt.Close()

	for _, file := range files {
		if file.IsDir() {
			show_path := filepath.Join(root, file.Name())
			if insert_tvshow_nfo(show_path, stmt) {
				err = scan_tvshow_episodes(show_path, tx)
				if err != nil {
					fmt.Fprintf(os.Stderr, "Could not scan for episodes in %v: %v", show_path, err)
					continue
				}
			}
		}
	}
	return nil
}

func insert_tvshow_nfo(show_dir string, stmt *sql.Stmt) bool {
	tvshow_nfo_path := filepath.Join(show_dir, "tvshow.nfo")

	tvshow_nfo_file, err := os.OpenFile(tvshow_nfo_path, os.O_RDONLY, 0755)
	if err != nil {
		log.Debugf("Could not open tvshow.nfo file in \"%v\"\n", show_dir)
		return false
	}
	defer tvshow_nfo_file.Close()

	// try to parse nfo file
	data, err := ioutil.ReadAll(tvshow_nfo_file)
	if err != nil {
		log.Debugf("Could not read data from tvshow.nfo file: %v\n", err)
		return false
	}
	var tvshow TVShow

	err = xml.Unmarshal(data, &tvshow)
	if err != nil {
		log.Debugf("Could not parse tvshow.nfo file contents: %v\n", err)
		return false
	}

	// TODO: Replace with print formatting logic
	log.Debugf("%v (%v):\n", tvshow.Title, tvshow.Year)
	log.Debugf("    Original Title: %v\n", tvshow.OriginalTile)
	log.Debugf("    Show Title: %v\n", tvshow.ShowTitle)

	_, err = stmt.Exec(tvshow.Title, tvshow.OriginalTile, tvshow.ShowTitle, tvshow.Year)
	if err != nil {
		log.Debugf("Could not insert tvshow \"%v\" into database: %v\n", tvshow.Title, err)
		return false
	}
	return true
}

type Episode struct {
	XMLName          xml.Name `xml:"episodedetails" json:"-"`
	Title            string   `xml:"title"`
	ShowTitle        string   `xml:"showtitle"`
	Season           int      `xml:"season"`
	Episode          int      `xml:"episode"`
	OriginalFilename string   `xml:"original_filename"`
}

func scan_tvshow_episodes(show_dir string, tx *sql.Tx) error {

	log.Infof("Scanning episodes for show %v\n", show_dir)

	files, err := ioutil.ReadDir(show_dir)
	if err != nil {
		return fmt.Errorf("Could not enumerate files in folder: %v", err)
	}

	num_episodes := 0

	stmt, err := tx.Prepare("insert into episodes(title, show_title, number, season, original_filename) values(?, ?, ?, ?, ?)")
	if err != nil {
		return fmt.Errorf("Could not prepare statement: %v", err)
	}
	defer stmt.Close()

	for _, file := range files {
		ext := filepath.Ext(file.Name())
		nfo_path := filepath.Join(show_dir, file.Name())

		if file.Name() == "tvshow.nfo" {
			continue
		} else if ext == ".nfo" {
			log.Debugf("Found nfo file: \"%v\"\n", file.Name())
		} else {
			log.Debugf("Skipping file \"%v\", not an nfo file\n", file.Name())
			continue
		}

		log.Debugf("%v\n", file.Name())

		nfo_file, err := os.OpenFile(nfo_path, os.O_RDONLY, 0755)
		if err != nil {
			log.Debugf("Could not open nfo file \"%v\": %v", nfo_path, err)
			continue
		}
		defer nfo_file.Close()

		// try to parse nfo file
		data, err := ioutil.ReadAll(nfo_file)
		if err != nil {
			log.Debugf("Could not read data from nfo file \"%v\": %v\n", nfo_path, err)
			continue
		}
		var episode Episode

		err = xml.Unmarshal(data, &episode)
		if err != nil {
			log.Debugf("Could not parse contents of nfo file \"%v\": %v\n", nfo_path, err)
			continue
		}

		log.Debugf("    Title: %v\n", episode.Title)
		log.Debugf("    Season %v\n", episode.Season)
		log.Debugf("    Episode %v\n", episode.Episode)
		log.Debugf("    Show: %v\n", episode.ShowTitle)
		log.Debugf("    Original Filename: %v\n", episode.OriginalFilename)

		full_original_filename := filepath.Join(show_dir, episode.OriginalFilename)

		_, err = stmt.Exec(episode.Title, episode.ShowTitle, episode.Episode, episode.Season, full_original_filename)
		if err != nil {
			log.Debugf("Could not insert episode %v of tvshow \"%v\" into database: %v\n", episode.Episode, episode.ShowTitle, err)
			continue
		}

		num_episodes++

	}

	log.Infof("Found %v episodes\n", num_episodes)

	return nil
}


type GetShows struct {
	Shows []TVShow
}
func get_shows_handler(data_dir string, w http.ResponseWriter, r *http.Request) {
	log.Debugf("/get_shows called\n")
	db, err := open_database(data_dir)
	if err != nil {
		log.Errorf("Could not open database: %v\n", err)
		return
	}

	rows, err := db.Query("SELECT title, original_title, show_title, year FROM tvshows")
	if err != nil {
		log.Errorf("Could not check if database was new: %v", err)
	}
	defer rows.Close()

	var return_json GetShows

	for {
		if !rows.Next() {
			err = rows.Err()
			if err != nil {
				log.Fatalf("Could not get row from database: %v\n", err)
				return
			} else {
				log.Debugf("%v rows processed\n", len(return_json.Shows))
				break
			}
		}
		var show TVShow
		rows.Scan(&show.Title, &show.OriginalTile, &show.ShowTitle, &show.Year)

		return_json.Shows = append(return_json.Shows, show)
	}
	data, err := json.Marshal(return_json)
	if err != nil {
		log.Errorf("Could not marshal episode data: %v\n", err)
		return
	}
	w.Write(data)
}

type GetEpisodes struct {
	Episodes []Episode
}

func get_episodes_handler(data_dir string, w http.ResponseWriter, r *http.Request) {
	log.Debugf("/get_episodes called\n")
	db, err := open_database(data_dir)
	if err != nil {
		log.Errorf("Could not open database: %v\n", err)
		return
	}

	rows, err := db.Query("SELECT title, show_title, number, season, original_filename FROM episodes")
	if err != nil {
		log.Errorf("Could not check if database was new: %v", err)
	}
	defer rows.Close()

	var return_json GetEpisodes

	for {
		if !rows.Next() {
			err = rows.Err()
			if err != nil {
				log.Fatalf("Could not get row from database: %v\n", err)
				return
			} else {
				log.Debugf("%v rows processed\n", len(return_json.Episodes))
				break
			}
		}
		var episode Episode
		rows.Scan(&episode.Title, &episode.ShowTitle, &episode.Episode, &episode.Season, &episode.OriginalFilename)

		return_json.Episodes = append(return_json.Episodes, episode)
	}
	data, err := json.Marshal(return_json)
	if err != nil {
		log.Errorf("Could not marshal episode data: %v\n", err)
		return
	}
	w.Write(data)
}

var log = logging.MustGetLogger("ikinuki-server")

func main() {
	var format = logging.MustStringFormatter(
		`%{color}%{time:15:04:05.000} %{level:.4s}:%{color:reset} %{message}`,
	)

	backend := logging.NewLogBackend(os.Stderr, "", 0)
	backendFormatter := logging.NewBackendFormatter(backend, format)
	backendLeveled := logging.AddModuleLevel(backendFormatter)
	backendLeveled.SetLevel(logging.DEBUG, "")

	logging.SetBackend(backendLeveled)

	config_dir := get_config_dir()

	config, err := get_config(config_dir)
	if err != nil {
		panic(err)
	}
	if config == nil {
		log.Warningf("Created config file at %v, please edit and launch again\n", config_dir)
		os.Exit(0)
	}
	log.Debugf("Config: %v\n", config)

	data_dir := get_data_dir()

	db, err := open_database(data_dir)
	if err != nil {
		log.Fatalf("Could not open database: %v", err)
	}
	defer db.Close()

	err = scan_tvshow_root(config.ScanDirectory, db)
	if err != nil {
		log.Fatalf("Could not scan directory \"%v\": %v\n", config.ScanDirectory, err)
	}

	http.HandleFunc("/get_episodes", func(w http.ResponseWriter, r *http.Request) {
		get_episodes_handler(data_dir, w, r)
	})
	http.HandleFunc("/get_shows", func(w http.ResponseWriter, r *http.Request) {
		get_shows_handler(data_dir, w, r)
	})


	connect_string := fmt.Sprintf("%v:%v", config.Address, config.Port)
	log.Noticef("Launching HTTP server on %v\n", connect_string)
	log.Fatal(http.ListenAndServe(connect_string, nil))
}