#!/usr/bin/env python3

from inscriptis import annotation
import inscriptis, re, requests, sys
from bs4 import BeautifulSoup
from glob import glob

name = sys.argv[1] or sys.exit(1)

def reorder(text, width):
    chunks = text.split("\n\n")
    out = []
    for i,c in enumerate(chunks):
        if (":" in c or "* " in c) and "Follow-up" not in c:
            for l in c.split("\n"):
                prefix = len(l) - len(l.replace("-", " ").lstrip())
                indent = l[:prefix]
                if i > 0 and chunks[i-1].startswith("*Example "):
                    indent = "  " + indent
                l2 = []
                for w in l[prefix:].split(" "):
                    if len(" ".join(l2 + [w])) > width:
                        out.append(indent + " ".join(l2))
                        indent = indent.replace("-", " ")
                        l2 = []
                    l2.append(w)
                if l == "" or l2 != []:
                    out.append(indent + " ".join(l2))
            out.append("")
            continue
        words = re.split(" |\n|\t", c)
        line = ""
        for w in words:
            if w == "":
                continue
            if len(line) + len(w) + 1 > width:
                out.append(line)
                while len(w) > width:
                    out.append(w[:width])
                    w = w[80:]
                line = ""
            line += "" if line == "" else " "
            line += w;
        if i == len(chunks) - 1:
            out.append(line)
        else:
            out.append(line)
            out.append("")
    return out

json_data = {
    'query': '\n    query questionContent($titleSlug: String!) {\n  question(titleSlug: $titleSlug) {\n    content\n    mysqlSchemas\n    dataSchemas\n  }\n}\n    ',
    'variables': {
        'titleSlug': name,
    },
    'operationName': 'questionContent',
}
r = requests.post('https://leetcode.com/graphql/', json=json_data)
description_html = r.json()["data"]["question"]["content"]
description_info = inscriptis.get_annotated_text(description_html,
    inscriptis.ParserConfig(annotation_rules = {"strong": ["emphasis"]}))
emphasis = sorted(description_info["label"], key=lambda v: v[0])
description = description_info["text"]
offset = 0
for start,end,_ in emphasis:
    inner = description[start+offset:end+offset]
    left = len(inner) - len(inner.lstrip())
    description = description[:start+offset+left] \
        + "*" + inner[left:] \
        + "*" + description[end+offset:]
    offset += 2
lines = description.split("\n")
plen = min([len(l) - len(l.lstrip()) for l in lines if l.strip() != ""])
lines = [l[plen:].rstrip() for l in lines]
lines = re.sub("\n\n+", "\n\n", "\n".join(lines)).split("\n")
lines = reorder("\n".join(lines), 76)
#lines = "\n".join(lines).strip().split("\n")
description_comment = "\n".join(["/// " + l if l != "" else "///" for l in lines])

json_data = {
    'query': '\n    query questionTitle($titleSlug: String!) {\n  question(titleSlug: $titleSlug) {\n    questionId\n    questionFrontendId\n    title\n    titleSlug\n    isPaidOnly\n    difficulty\n    likes\n    dislikes\n    categoryTitle\n  }\n}\n    ',
    'variables': {
        'titleSlug': name,
    },
    'operationName': 'questionTitle',
}
r = requests.post('https://leetcode.com/graphql/', json=json_data)
num = int(r.json()["data"]["question"]["questionFrontendId"])
title = r.json()["data"]["question"]["title"]
difficulty = r.json()["data"]["question"]["difficulty"]
info_comment = f"/// {num}. {title} ({difficulty})\n///"
header_comment = info_comment + "\n" + description_comment

json_data = {
    'query': '\n    query questionEditorData($titleSlug: String!) {\n  question(titleSlug: $titleSlug) {\n    questionId\n    questionFrontendId\n    codeSnippets {\n      lang\n      langSlug\n      code\n    }\n    envInfo\n    enableRunCode\n    hasFrontendPreview\n    frontendPreviews\n  }\n}\n    ',
    'variables': {
        'titleSlug': name,
    },
    'operationName': 'questionEditorData',
}
r = requests.post('https://leetcode.com/graphql/', json=json_data)
solution_impl = next(c for c in r.json()["data"]["question"]["codeSnippets"] if c["lang"] == "Rust")["code"]

main_footer = """
#[cfg(test)]
mod tests {
    use crate::Solution;

    #[test]
    fn test() {
    }
}
"""
with open(f"src/bin/{num:04}-{name}.rs", "w+") as f:
    f.write(header_comment + "\n\n")
    f.write("use leetcode::*;\n\n")
    f.write("struct Solution {}\n\n")
    f.write(solution_impl + "\n\n")
    f.write("pub fn main() {\n}\n")
    f.write(main_footer)

cargo_header = """[package]
name = "leetcode"
version = "0.1.0"
edition = "2021"
"""
with open("Cargo.toml", "w+") as f:
    f.write(cargo_header)
    f.write("\n")
    for path in glob("src/bin/*.rs"):
        num = int(path.rsplit("/",1)[1].split("-",1)[0])
        f.write("[[bin]]\n")
        f.write(f"name = \"{num}\"\n")
        f.write(f"path = \"{path}\"\n")
        f.write("\n")