import argparse
import collections
import http.server
import json
import os
import pathlib
import re
import shutil
import jinja2

from markdown.core import Markdown
from markdown.extensions import Extension
from markdown.preprocessors import Preprocessor
from obsidian_parser import Note, Vault


def out_content(self):
    content = self.content
    if content.startswith("---\n"):
        content = content.split("---\n", 2)[2]
    return content
Note.our_content = property(out_content)

root_dir = pathlib.Path(__file__).parent
vault = Vault(root_dir.parent / "content")
loader = jinja2.FileSystemLoader(root_dir / "templates")
env = jinja2.Environment(loader=loader)


class ObsidianImage(Preprocessor):
    def run(self, lines):
        new_lines = []
        for line in lines:
            m = re.search(r"!\[\[(.*)\]\]", line)
            if m:
                if "|" in m.group(1):
                    img, width = m.group(1).split("|")
                    new_lines.append(f"<img src='attachments/{img}' style='width:{width};'/>")
                else:
                    new_lines.append(f"<img src='attachments/{m.group(1)}'/>")
            else:
                new_lines.append(line)
        return new_lines


class ObsidianImageExtension(Extension):
    def extendMarkdown(self, md):
        md.preprocessors.register(ObsidianImage(md), 'obsidianimage', 175)

def make_markdown():
    return Markdown(
        extensions=[
            "fenced_code",
            "mdx_math",
            "nl2br",
            "tables",
            ObsidianImageExtension(),
        ],
        extension_configs={
            "mdx_math": {
                "enable_dollar_delimiter": True
            }
        },
        tab_length=2,
    )

def markdown_filter(text):
    md = make_markdown()
    return md.convert(text)

env.filters["markdown"] = markdown_filter
output_dir = root_dir / "output"


def build_tree(vault: Vault):
    root = vault.path
    tree = {}
    for folder, dirnames, filenames in root.walk():
        for dirname in dirnames[:]:
            if dirname.startswith(".") or dirname == "attachments":
                dirnames.remove(dirname)
        cur = tree
        for part in folder.relative_to(root).parts:
            cur = cur.setdefault(part, {"children": {}})["children"]

        for filename in filenames:
            # .DS_Store etc
            if filename.startswith("."):
                continue
            item = {
                "filename": filename,
                "folder": str(folder.relative_to(root)),
            }
            note = vault.get_note((folder / filename).relative_to(root))
            if note:
                item["title"] = note.title
                # TODO: for search add tags, modified time, content etc
            cur[filename] = item
    return tree


def write_note(item):
    if "children" in item:
        for child in item["children"].values():
            write_note(child)
    else:
        path = pathlib.Path(item["folder"]) / item["filename"]
        context = {
            "filename": item["filename"],
            "folder": item["folder"],
            "path": path,
        }

        folder = output_dir / item["folder"]
        folder.mkdir(parents=True, exist_ok=True)
        link = False
        if note := vault.get_note(path):
            template_name = "note.jinja2"
            context["note"] = note
            context["title"] = note.title
        elif item["filename"].endswith(".pdf"):
            template_name = "pdf.jinja2"
            context["title"] = item["filename"]
            link = True
        elif item["filename"].lower().endswith((".png", ".jpg", ".jpeg", ".gif", ".webp")):
            template_name = "image.jinja2"
            context["title"] = item["filename"]
            link = True
        else:
            print(f"Note not found for {path}")
            return

        if link:
            os.link(
                vault.path / item["folder"] / item["filename"],
                output_dir / item["folder"] / item["filename"],
            )

        out_path = folder / (context["title"] + ".html")
        with out_path.open("w", encoding="utf-8") as f:
            template = env.get_template(template_name)
            data = template.render(**context)
            f.write(data)


def write_json(output_path, tree):
    tree_json = json.dumps(tree)
    with output_path.open("w") as f:
        f.write(tree_json)

def build():
    """Build the static site."""
    print("Building...")

    shutil.rmtree(output_dir, ignore_errors=True)
    output_dir.mkdir(exist_ok=True)

    (output_dir / "style.css").symlink_to(root_dir / "style.css")
    (output_dir / "script.js").symlink_to(root_dir / "script.js")

    attachments_src = root_dir.parent / "content" / "attachments"
    attachments_dst = output_dir / "attachments"
    attachments_dst.symlink_to(attachments_src)

    tree = build_tree(vault)
    write_json(output_dir / "index.json", tree)

    # write quiz
    quiz = []
    for note in vault.get_notes_with_tag("provfråga"):
        # assuming structure content/exams/YYYY-MM-DD/Note.md
        exam_date = note.path.parent.parent.name
        quiz.append({
            "question": note.title,
            "date": exam_date,
            "number": note.path.name,
        })
    write_json(output_dir / "quiz.json", quiz)


    # 3. Write out each note as html
    write_note({"children": tree})

    print(f"Built to {output_dir}")


def serve(port):
    """Build and serve the site."""
    build()

    os.chdir(output_dir)
    print(f"Serving at http://localhost:{port}")
    server = http.server.HTTPServer(("", port), http.server.SimpleHTTPRequestHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped.")


def main():
    parser = argparse.ArgumentParser(description="Static site generator for medical notes")
    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser("build", help="Build the static site")

    serve_parser = subparsers.add_parser("serve", help="Build and serve the site")
    serve_parser.add_argument("-p", "--port", type=int, default=8000, help="Port to serve on")

    args = parser.parse_args()

    if args.command == "build":
        build()
    elif args.command == "serve":
        serve(args.port)


if __name__ == "__main__":
    main()