import argparse
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 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)
    tree_json = json.dumps(tree)
    with (output_dir / "index.json").open("w") as f:
        f.write(tree_json)

    # 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()