package main

import (
	"compress/gzip"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
)

var (
	docType = `<!DOCTYPE tv SYSTEM "xmltv.dtd">`

	epgRE     = regexp.MustCompile(`url-tvg="([a-z0-9-:/.]+)"`)
	channelRE = regexp.MustCompile(`tvg-id="([a-z0-9-]+)"`)
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: epgpruner URL\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		usage()
		os.Exit(1)
	}
	playlistURL := flag.Args()[0]
	playlist := Playlist{
		URL: playlistURL,
	}
	err := playlist.Fetch()
	if err != nil {
		log.Fatal(err)
	}

	playlist.Reduce()

	reducedEpg, err := xml.MarshalIndent(playlist.Epg, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(
		xml.Header +
			docType + "\n" +
			string(reducedEpg),
	)
}

type Playlist struct {
	URL      string
	EpgURL   string
	Channels map[string]struct{}
	Epg      Epg
}

func (p *Playlist) Fetch() error {
	compressed := false
	body, err := fetch(p.URL, compressed)
	if err != nil {
		return err
	}
	err = p.parse(body)
	if err != nil {
		return err
	}

	compressed = true
	body, err = fetch(p.EpgURL, compressed)
	if err != nil {
		return err
	}

	err = xml.Unmarshal(body, &p.Epg)
	if err != nil {
		return err
	}

	return nil
}

func fetch(url string, compressed bool) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return make([]byte, 0), err
	}
	if resp.StatusCode != http.StatusOK {
		return make([]byte, 0),
			fmt.Errorf("Bad reponse status: %d", resp.StatusCode)
	}
	reader := resp.Body
	defer reader.Close()
	if compressed {
		reader, err = gzip.NewReader(reader)
		if err != nil {
			return make([]byte, 0), err
		}
	}
	body, err := io.ReadAll(reader)
	if err != nil {
		return make([]byte, 0), err
	}
	return body, nil
}

func (p *Playlist) parse(body []byte) error {
	epgURL := ""
	channels := make(map[string]struct{})
	for _, line := range strings.Split(string(body), "\n") {
		if strings.HasPrefix(line, "#EXTM3U") {
			m := epgRE.FindStringSubmatch(line)
			if m == nil {
				return fmt.Errorf(
					"Couldn't match EPG url from: `%s`",
					line,
				)
			}
			epgURL = m[1]
		} else if strings.HasPrefix(line, "#EXTINF") {
			m := channelRE.FindStringSubmatch(line)
			if m == nil {
				return fmt.Errorf(
					"Couldn't match channel name from: `%s`",
					line,
				)
			}
			channel := m[1]
			channels[channel] = struct{}{}
		}
	}

	p.EpgURL = epgURL
	p.Channels = channels

	return nil
}

func (p *Playlist) Reduce() {
	channels := make([]Channel, 0)
	programmes := make([]Programme, 0)

	for _, channel := range p.Epg.Channels {
		if _, ok := p.Channels[channel.ID]; !ok {
			continue
		}
		channels = append(channels, channel)
	}

	for _, programme := range p.Epg.Programmes {
		if _, ok := p.Channels[programme.Channel]; !ok {
			continue
		}
		programmes = append(programmes, programme)
	}

	p.Epg.Channels = channels
	p.Epg.Programmes = programmes
}

type Epg struct {
	XMLName    xml.Name    `xml:"tv"`
	Info       string      `xml:"generator-info-name,attr"`
	Channels   []Channel   `xml:"channel"`
	Programmes []Programme `xml:"programme"`
}

type Channel struct {
	ID          string      `xml:"id,attr"`
	DisplayName DisplayName `xml:"display-name"`
	Icon        Icon        `xml:"icon"`
}

type DisplayName struct {
	Lang string `xml:"lang,attr"`
	Name string `xml:",chardata"`
}

type Icon struct {
	Src string `xml:"src,attr"`
}

type Programme struct {
	Start   string `xml:"start,attr"`
	Stop    string `xml:"stop,attr"`
	Channel string `xml:"channel,attr"`
	Title   Title  `xml:"title,omitempty"`
	Desc    Desc   `xml:"desc,omitempty"`
}

type Title struct {
	Lang  string `xml:"lang,attr"`
	Title string `xml:",chardata"`
}

type Desc struct {
	Lang string `xml:"lang,attr"`
	Desc string `xml:",chardata"`
}