package main

import (
	"bufio"
	"compress/gzip"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
)

type Playlist struct {
	URL      string
	Channels map[string]struct{}
	Schedule Schedule
}

type Schedule struct {
	URL string
	EPG EPG
}

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"`
}

func NewPlaylist(URL string) (*Playlist, error) {
	pr, err := open(URL)
	if err != nil {
		return nil, err
	}
	defer pr.Close()

	scheduleURL, channels, err := parsePlaylist(pr)
	if err != nil {
		return nil, err
	}

	sr, err := open(scheduleURL)
	if err != nil {
		return nil, err
	}
	csr, err := gzip.NewReader(sr)
	if err != nil {
		return nil, err
	}
	epg, err := parseEPG(csr, channels)
	if err != nil {
		return nil, err
	}
	defer sr.Close()

	playlist := Playlist{
		URL:      URL,
		Channels: channels,
		Schedule: Schedule{
			URL: scheduleURL,
			EPG: epg,
		},
	}
	return &playlist, nil
}

func open(URL string) (io.ReadCloser, error) {
	resp, err := http.Get(URL)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil,
			fmt.Errorf("Bad reponse status: %d", resp.StatusCode)
	}
	return resp.Body, nil
}

func parsePlaylist(r io.Reader) (string, map[string]struct{}, error) {
	scheduleURL := ""
	channels := make(map[string]struct{})

	br := bufio.NewReader(r)

	for {
		line, err := br.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				return "", nil, err
			}
			break
		}

		if strings.HasPrefix(line, "#EXTM3U") {
			scheduleURL, err = parseScheduleURL(line)
			if err != nil {
				return "", nil, err
			}
		} else if strings.HasPrefix(line, "#EXTINF") {
			channelID, err := parseChannelID(line)
			if err != nil {
				return "", nil, err
			}
			channels[channelID] = struct{}{}
		}
	}

	return scheduleURL, channels, nil
}

var scheduleRE = regexp.MustCompile(`url-tvg="([a-z0-9-:/.]+)"`)

func parseScheduleURL(line string) (string, error) {
	m := scheduleRE.FindStringSubmatch(line)
	if m == nil {
		return "", fmt.Errorf(
			"Couldn't match schedule URL from: `%s`",
			line,
		)
	}
	return m[1], nil
}

var channelRE = regexp.MustCompile(`tvg-id="([a-z0-9-]+)"`)

func parseChannelID(line string) (string, error) {
	m := channelRE.FindStringSubmatch(line)
	if m == nil {
		return "", fmt.Errorf(
			"Couldn't match channel name from: `%s`",
			line,
		)
	}
	return m[1], nil
}

func parseEPG(r io.Reader, channels map[string]struct{}) (EPG, error) {
	d := xml.NewDecoder(r)

	var epg EPG

	for {
		t, err := d.Token()
		if err != nil {
			if err != io.EOF {
				return epg, err
			}
			break
		}

		switch et := t.(type) {
		case xml.StartElement:
			if et.Name.Local == "tv" {
				for _, attr := range et.Attr {
					name := attr.Name.Local
					if name != "generator-info-name" {
						continue
					}
					epg.Info = attr.Value
					break
				}
			} else if et.Name.Local == "channel" {
				// We can check if attribute value is
				// in the channels map without decoding
				// entire element, but I am too lazy
				// for that. Channel have only couple of
				// fields, so this should be good enough.
				var channel Channel
				err := d.DecodeElement(&channel, &et)
				if err != nil {
					return epg, err
				}
				if _, ok := channels[channel.ID]; !ok {
					continue
				}
				epg.Channels = append(
					epg.Channels, channel,
				)
			} else if et.Name.Local == "programme" {
				var programme Programme
				err := d.DecodeElement(&programme, &et)
				if err != nil {
					return epg, err
				}
				channel := programme.Channel
				if _, ok := channels[channel]; !ok {
					continue
				}
				epg.Programmes = append(
					epg.Programmes, programme,
				)
			}
		}
	}

	return epg, nil
}