From c15d7607eb81418f6f2b73a0c67896a13368db39 Mon Sep 17 00:00:00 2001
From: Harald Gustafsson <harald.gustafsson@gmail.com>
Date: Sun, 22 Jul 2018 11:17:04 +0200
Subject: [PATCH] Create NFO files with the extracted metadata. TVshows get a
 tvshow.nfo file as well.

---
 lib/svtplay_dl/utils/getmedia.py |  6 ++++
 lib/svtplay_dl/utils/nfo.py      | 47 ++++++++++++++++++++++++++++++++
 lib/svtplay_dl/utils/parser.py   |  4 +++
 3 files changed, 57 insertions(+)
 create mode 100644 lib/svtplay_dl/utils/nfo.py

diff --git a/lib/svtplay_dl/utils/getmedia.py b/lib/svtplay_dl/utils/getmedia.py
index 0825221..7ead5d8 100644
--- a/lib/svtplay_dl/utils/getmedia.py
+++ b/lib/svtplay_dl/utils/getmedia.py
@@ -15,6 +15,7 @@ from svtplay_dl.postprocess import postprocess
 from svtplay_dl.utils.stream import select_quality, list_quality
 from svtplay_dl.utils.text import exclude
 from svtplay_dl.error import UIException
+from svtplay_dl.utils.nfo import write_nfo_episode, write_nfo_tvshow
 
 
 def get_multiple_media(urls, config):
@@ -188,6 +189,11 @@ def get_one_media(stream):
 
         if fstream.config.get("thumbnail") and hasattr(stream, "get_thumbnail"):
             stream.get_thumbnail(stream.config)
+        if stream.config.get("nfo"):
+            # Create NFO files
+            write_nfo_episode(stream.output, stream.config)
+            write_nfo_tvshow(stream.output, stream.config)
+
         post = postprocess(fstream, fstream.config, subfixes)
         if fstream.audio and post.detect:
             post.merge()
diff --git a/lib/svtplay_dl/utils/nfo.py b/lib/svtplay_dl/utils/nfo.py
new file mode 100644
index 0000000..5f38db7
--- /dev/null
+++ b/lib/svtplay_dl/utils/nfo.py
@@ -0,0 +1,47 @@
+import logging
+
+import xml.etree.ElementTree as ET
+from svtplay_dl.utils.output import formatname
+from svtplay_dl.utils.parser import Options
+
+
+def write_nfo_episode(output, config):
+    root = ET.Element("episodedetails")
+    ET.SubElement(root, "title").text = output["title_nice"]
+    ET.SubElement(root, "showtitle").text = output["episodename"]
+    ET.SubElement(root, "season").text = output["season"]
+    ET.SubElement(root, "episode").text = output["episode"]
+    ET.SubElement(root, "plot").text = output["showdescription"]
+    if not config.get("thumbnail"):
+        # Set the thumbnail path to download link if not thumbnail downloaded
+        ET.SubElement(root, "thumb").text = output["showthumbnailurl"]
+
+    filename = formatname(output.copy(), config, extension="nfo")
+    logging.info("NFO episode: %s", filename)
+
+    tree = ET.ElementTree(root)
+    tree.write(filename, encoding="UTF-8", xml_declaration=True)
+
+
+def write_nfo_tvshow(output, config):
+    # Config for tvshow nfo file
+    root = ET.Element("tvshow")
+    ET.SubElement(root, "title").text = output["title_nice"]
+    ET.SubElement(root, "plot").text = output["episodedescription"]
+    if config.get("thumbnail"):
+        # Set the thumbnail relative path to downloaded thumbnail
+        ET.SubElement(root, "thumb").text = "{}.tvshow.tbn".format(output["title"])
+    else:
+        # Set the thumbnail path to download link if not thumbnail downloaded
+        ET.SubElement(root, "thumb").text = output["episodethumbnailurl"]
+
+    cconfig = Options()
+    cconfig.set("output", config.get("output"))
+    cconfig.set("path", config.get("path"))
+    cconfig.set("subfolder", config.get("subfolder"))
+    cconfig.set("filename", "tvshow.{ext}")
+    filename = formatname(output.copy(), cconfig, extension="nfo")
+    logging.info("NFO show: %s", filename)
+
+    tree = ET.ElementTree(root)
+    tree.write(filename, encoding="UTF-8", xml_declaration=True)
\ No newline at end of file
diff --git a/lib/svtplay_dl/utils/parser.py b/lib/svtplay_dl/utils/parser.py
index f3c992c..9968afd 100644
--- a/lib/svtplay_dl/utils/parser.py
+++ b/lib/svtplay_dl/utils/parser.py
@@ -92,6 +92,8 @@ def parser(version):
                                                "socks5://127.0.0.1:1080/.")
     general.add_argument("-v", "--verbose", action="store_true", dest="verbose", default=False,
                          help="explain what is going on")
+    general.add_argument("--nfo", action="store_true", dest="nfo", default=False,
+                         help="create a NFO file")
 
     quality = parser.add_argument_group("Quality")
     quality.add_argument("-q", "--quality", default=0, metavar="quality",
@@ -171,6 +173,7 @@ def setup_defaults():
     options.set("convert_subtitle_colors", False)
     options.set("preferred", None)
     options.set("verbose", False)
+    options.set("nfo", False)
     options.set("output_auto", False)
     options.set("service", None)
     options.set("cookies", None)
@@ -212,6 +215,7 @@ def parsertoconfig(config, parser):
     config.set("require_subtitle", parser.require_subtitle)
     config.set("preferred", parser.preferred)
     config.set("verbose", parser.verbose)
+    config.set("nfo", parser.nfo)
     config.set("exclude", parser.exclude)
     config.set("get_url", parser.get_url)
     config.set("ssl_verify", parser.ssl_verify)