#!/usr/bin/env python3
import sys
import re
from pathlib import Path

INDENT = " " * 4

CLASS_START_RE = re.compile(
    r"^\s*(?:partial\s+)?(?:package|model|connector|record|function|block|type|operator)\b"
)

SECTION_RE = re.compile(
    r"^\s*(?:equation|algorithm|initial\s+equation|initial\s+algorithm|protected|public)\b"
)

IF_RE = re.compile(r"^\s*if\b.*\bthen\b")
ELSEIF_RE = re.compile(r"^\s*elseif\b.*\bthen\b")
ELSE_RE = re.compile(r"^\s*else\b")
FOR_RE = re.compile(r"^\s*for\b.*\bloop\b")
WHEN_RE = re.compile(r"^\s*when\b.*\bthen\b")
ELSEWHEN_RE = re.compile(r"^\s*elsewhen\b.*\bthen\b")

END_RE = re.compile(r"^\s*end\b")


def format_modelica(text: str) -> str:
    lines = text.replace("\t", "    ").splitlines()
    out = []
    indent = 0

    for raw in lines:
        stripped = raw.strip()

        if not stripped:
            out.append("")
            continue

        # Closing constructs outdent before printing
        if END_RE.match(stripped):
            indent = max(indent - 1, 0)

        # else / elseif / elsewhen align with matching if/when
        elif ELSE_RE.match(stripped) or ELSEIF_RE.match(stripped) or ELSEWHEN_RE.match(stripped):
            indent = max(indent - 1, 0)

        # Sections stay at model level, not one level deeper
        elif SECTION_RE.match(stripped):
            indent = max(indent - 1, 0)

        out.append(f"{INDENT * indent}{stripped}")

        # Increase indent only for real blocks
        if CLASS_START_RE.match(stripped):
            indent += 1
        elif IF_RE.match(stripped):
            indent += 1
        elif FOR_RE.match(stripped):
            indent += 1
        elif WHEN_RE.match(stripped):
            indent += 1
        elif ELSEIF_RE.match(stripped) or ELSE_RE.match(stripped) or ELSEWHEN_RE.match(stripped):
            indent += 1
        elif SECTION_RE.match(stripped):
            indent += 1

    return "\n".join(out) + "\n"


def main():
    if len(sys.argv) < 2:
        print("Usage: python format_mo.py input.mo [output.mo]")
        sys.exit(1)

    input_file = Path(sys.argv[1])
    output_file = Path(sys.argv[2]) if len(sys.argv) > 2 else input_file

    text = input_file.read_text(encoding="utf-8")
    formatted = format_modelica(text)
    output_file.write_text(formatted, encoding="utf-8")

    print(f"Formatted {output_file}")


if __name__ == "__main__":
    main()