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

template = env.get_template("base.html")
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, tree_json):
    if "children" in item:
        for child in item["children"].values():
            write_note(child, tree_json)
    else:
        path = pathlib.Path(item["folder"]) / item["filename"]
        note = vault.get_note(path)
        if note:
            out_path = output_dir / item["folder"] / (item["title"] + ".html")
            out_path.parent.mkdir(parents=True, exist_ok=True)

            # Calculate relative base_path based on folder depth
            folder = item["folder"]
            if folder == ".":
                base_path = ""
            else:
                depth = len(pathlib.Path(folder).parts)
                base_path = "../" * depth

            with out_path.open("w", encoding="utf-8") as f:
                data = template.render(note=note, vault=vault, base_path=base_path, index_json=tree_json)
                f.write(data)
        else:
            print(f"Note not found for {path}")


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

    # 1. Create output dir
    shutil.rmtree(output_dir, ignore_errors=True)
    output_dir.mkdir(exist_ok=True)

    # 1b. Symlink CSS/JS to output root
    (output_dir / "style.css").symlink_to(root_dir / "style.css")
    (output_dir / "script.js").symlink_to(root_dir / "script.js")

    # 1c. Symlink attachments directory
    attachments_src = root_dir.parent / "content" / "attachments"
    attachments_dst = output_dir / "attachments"
    attachments_dst.symlink_to(attachments_src)

    # 2. Build tree and write index json
    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}, tree_json)

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