#!/usr/bin/env python3
import os
import sys
import subprocess
from pathlib import Path
import shutil
import platform

# ---------------------------------------------------------------------
# CONFIG
# ---------------------------------------------------------------------
SRC_DIR = Path("src")
TEST_DIR = Path("tests")
OBJ_DIR = Path("obj")
BIN_DIR = Path("bin")
PICO_GENERATED = Path("generated")

TARGET = BIN_DIR / "sls"
TEST_TARGET = BIN_DIR / "sls_tests"


# Platform-specific settings
PICO_SDK_PATH = os.environ.get("PICO_SDK_PATH", Path.home() / "pico/pico-sdk")

# Unix gcc/clang flags
COMMON_FLAGS = ["-std=c99", "-Wall", "-Wextra", "-Werror", "-Iinclude", "-g"]
PICO_FLAGS = ["-std=c11", "-Wall", "-Wextra", "-Werror", "-Iinclude", "-g"]
TEST_FLAGS = ["-std=c99", "-Wall", "-Wextra", "-Wno-unused-function", "-Werror",
              "-Iinclude", "-g", "-O0"]

# macOS-specific flags (for cross-compilation if needed)
MACOS_FLAGS = ["-std=c99", "-Wall", "-Wextra", "-Werror", "-Iinclude", "-g",
               "-mmacosx-version-min=10.13"]

# Windows MSVC flags
MSVC_FLAGS = ["/std:c11", "/Zi", "/Iinclude"]
MSVC_TEST_FLAGS = MSVC_FLAGS + []

# ---------------------------------------------------------------------
# RP2040 SETTINGS
# ---------------------------------------------------------------------
PICO_CPU_FLAGS = [
    "-mcpu=cortex-m0plus",
    "-mthumb",
    "-O2",
    "-ffunction-sections",
    "-fdata-sections"
]

PICO_DEFINES = [
    "-DPICO_BUILD=1",
    "-DPICO_ON_DEVICE=1",
]

PICO_ASM_INCLUDES = [
    PICO_GENERATED / "include",
    f"{PICO_SDK_PATH}/src/rp2040/boot_stage2/asminclude",
    f"{PICO_SDK_PATH}/src/rp2040/pico_platform/include",
    f"{PICO_SDK_PATH}/src/rp2040/hardware_regs/include",
    f"{PICO_SDK_PATH}/src/common/pico_base_headers/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_compiler/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_sections/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_panic/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_common/include",
    f"{PICO_SDK_PATH}/src/common/pico_binary_info/include",
    f"{PICO_SDK_PATH}/src/common/boot_picobin_headers/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_bootrom/include",
    f"{PICO_SDK_PATH}/src/rp2_common/boot_bootrom_headers/include",
]

PICO_INCLUDES = [
    PICO_GENERATED / "include",

    "src",
    "include",

    # Pico SDK core
    f"{PICO_SDK_PATH}/src/common/pico_stdlib_headers/include",
    f"{PICO_SDK_PATH}/src/common/pico_base_headers/include",
    f"{PICO_SDK_PATH}/src/common/pico_base_headers/include",
    f"{PICO_SDK_PATH}/src/common/pico_time/include",
    f"{PICO_SDK_PATH}/src/common/pico_sync/include",
    f"{PICO_SDK_PATH}/src/common/pico_binary_info/include",

    # RP2040 hardware + platform
    f"{PICO_SDK_PATH}/src/rp2040/hardware_regs/include",
    f"{PICO_SDK_PATH}/src/rp2040/hardware_structs/include",
    f"{PICO_SDK_PATH}/src/rp2040/pico_platform/include",

    # Runtime + stdio
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_compiler/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_sections/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_panic/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_platform_common/include",
    # f"{PICO_SDK_PATH}/src/rp2_common/pico_platform/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_runtime/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_stdio/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_stdio_usb/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_runtime_init/include",
    f"{PICO_SDK_PATH}/src/rp2_common/pico_stdio_uart/include",

    # Hardware libs you use
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_gpio/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_timer/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_base/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_gpio/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_irq/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_uart/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_resets/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_sync/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_sync_spin_lock/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_clocks/include",
    f"{PICO_SDK_PATH}/src/common/hardware_claim/include",
    f"{PICO_SDK_PATH}/src/rp2_common/hardware_pll/include",
]

PICO_LINKER_SCRIPT = (
    Path(PICO_SDK_PATH)
    / "src"
    / "rp2_common"
    / "pico_crt0"
    / "rp2040"
    / "memmap_default.ld"
)

PICO_LINKER_SCRIPTS_DIR = PICO_GENERATED / "link/pico"

def pico_sdk_asm_sources():
    sdk = Path(PICO_SDK_PATH)

    return [
        sdk / "src/rp2040/boot_stage2/boot2_generic_03h.S",
        sdk / "src/rp2_common/pico_crt0/crt0.S",
    ]

def compile_pico_asm_source(src: Path):
    mkdir(OBJ_DIR)
    obj = OBJ_DIR / (src.stem + ".pico.o")

    if obj.exists() and src.stat().st_mtime < obj.stat().st_mtime:
        return obj

    include_flags = [f"-I{p}" for p in PICO_ASM_INCLUDES]

    cmd = (
        ["arm-none-eabi-gcc"]
        + PICO_CPU_FLAGS
        + ["-DPICO_FLASH_SPI_CLKDIV=2"]
        + include_flags
        + ["-c", str(src), "-o", str(obj)]
    )

    run(cmd)
    return obj

def pico_sdk_sources():
    sdk = Path(PICO_SDK_PATH)

    return [
        sdk / "src/rp2_common/pico_runtime/runtime.c",
        sdk / "src/rp2_common/pico_runtime_init/runtime_init.c",
        sdk / "src/rp2040/pico_platform/platform.c",
        sdk / "src/rp2_common/pico_stdio/stdio.c",
        sdk / "src/rp2_common/pico_stdio_uart/stdio_uart.c",
        sdk / "src/rp2_common/hardware_uart/uart.c",
        sdk / "src/rp2_common/hardware_gpio/gpio.c",
        sdk / "src/rp2_common/pico_stdlib/stdlib.c",
        sdk / "src/rp2_common/hardware_clocks/clocks.c",
        sdk / "src/rp2_common/hardware_sync/sync.c",
        sdk / "src/rp2_common/hardware_irq/irq.c",
    ]

def compile_pico_source(src: Path):
    mkdir(OBJ_DIR)
    obj = OBJ_DIR / (src.stem + ".pico.o")

    if obj.exists() and src.stat().st_mtime < obj.stat().st_mtime:
        return obj

    include_flags = [f"-I{p}" for p in PICO_INCLUDES]

    cmd = (
        ["arm-none-eabi-gcc"]
        + PICO_CPU_FLAGS
        + PICO_FLAGS
        + PICO_DEFINES
        + include_flags
        + [f"-DGIT_COMMIT_HASH=\"{GIT_HASH}\""]
        + ["-c", str(src), "-o", str(obj)]
    )

    run(cmd)
    return obj

def link_pico_elf(objects, elf_path: Path):
    cmd = [
        "arm-none-eabi-gcc",
        "-nostdlib",
        "-Wl,--gc-sections",
        f"-Wl,-L{PICO_LINKER_SCRIPTS_DIR}",
        f"-T{PICO_LINKER_SCRIPT}",
        "-Wl,-Map=" + str(elf_path.with_suffix(".map")),
    ] + list(map(str, objects)) + [
        "-o",
        str(elf_path),
        "-lc",
        "-lgcc",
    ]

    run(cmd)

def elf_to_uf2(elf: Path):
    elf2uf2 = Path(PICO_SDK_PATH) / "tools/elf2uf2/elf2uf2"

    if not elf2uf2.exists():
        raise RuntimeError("elf2uf2 not found")

    uf2 = elf.with_suffix(".uf2")
    run([str(elf2uf2), str(elf), str(uf2)])
    return uf2

def generate_pico_headers():
    out = PICO_GENERATED / "include/pico"
    mkdir(out)

    version_template = Path(PICO_SDK_PATH) / "src/common/pico_base_headers/include/pico/version.h.in"
    version_out = out / "version.h"

    if not version_out.exists():

        version_out.write_text(version_template.read_text()
            .replace("${PICO_SDK_VERSION_MAJOR}", "1")
            .replace("${PICO_SDK_VERSION_MINOR}", "5")
            .replace("${PICO_SDK_VERSION_REVISION}", "1")
            .replace("${PICO_SDK_VERSION_STRING}", "1.5.1")
        )

    config_out = out / "config_autogen.h"

    if not config_out.exists():
        config_out.write_text("""// generated/pico/config_autogen.h
#pragma once

/* Platform */
#define PICO_RP2040 1

/* Board / runtime */
#define PICO_ON_DEVICE 1
#define PICO_NO_HARDWARE 0

/* stdio */
#define PICO_STDIO_UART 1
#define PICO_STDIO_USB 0

/* Assertions / debug */
#define PICO_ENABLE_ASSERTIONS 0
#define PICO_USE_STACK_GUARDS 0

/* Time */
#define PICO_TIME_DEFAULT_ALARM_POOL_DISABLED 1

/* Multicore (disable if unused) */
#define PICO_MULTICORE 0
""")

def generate_pico_linker_scripts():
    mkdir(PICO_LINKER_SCRIPTS_DIR)

    pico_flash_region_template = Path(PICO_SDK_PATH) / "src/rp2_common/pico_standard_link/pico_flash_region.template.ld"
    pico_flash_region_out = PICO_LINKER_SCRIPTS_DIR / "pico_flash_region.ld"

    if not pico_flash_region_out.exists():

        pico_flash_region_out.write_text(pico_flash_region_template.read_text()
            .replace("${PICO_FLASH_SIZE_BYTES_STRING}", hex(2 * 1024 * 1024))
        )

# ---------------------------------------------------------------------
# PLATFORM DETECTION
# ---------------------------------------------------------------------
def detect_platform():
    """Detect the current operating system"""
    system = platform.system()
    if system == "Darwin":
        return "macos"
    elif system == "Windows":
        return "windows"
    elif system == "Linux":
        return "linux"
    return "unknown"

PLATFORM = detect_platform()

# ---------------------------------------------------------------------
# COMPILER DETECTION
# ---------------------------------------------------------------------
def detect_compiler(target_platform=None):
    """Detect appropriate compiler for the target platform"""
    if target_platform == "rp2040":
        return ("arm-none-eabi-gcc", "gcc")
    
    target = target_platform or PLATFORM
    
    if target == "windows" or os.name == "nt":
        return ("cl", "msvc")
    elif target == "macos":
        # Prefer clang on macOS
        if shutil.which("clang"):
            return ("clang", "gcc")
        return ("gcc", "gcc")
    else:
        return ("gcc", "gcc")

CC, CC_KIND = detect_compiler()

# ---------------------------------------------------------------------
# PYTHON DETECTION
# ---------------------------------------------------------------------
def detect_python():
    if os.name == "nt":
        return "python"
    return "python3"

PYTHON = detect_python()


# ---------------------------------------------------------------------
# GIT COMMIT HASH
# ---------------------------------------------------------------------
def git_commit_hash():
    try:
        result_hash = subprocess.check_output(
            ["git", "describe", "--always", "--dirty", "--abbrev=7"],
            cwd=".",
            stderr=subprocess.DEVNULL,
            text=True
        ).strip()
        result_date = subprocess.check_output(
            ["git", "show", "-s", "--format=%ci"],
            cwd=".",
            stderr=subprocess.DEVNULL,
            text=True
        ).strip()
        return f"{result_hash} {result_date}"
    except Exception:
        return "unknown"


GIT_HASH = git_commit_hash()


# ---------------------------------------------------------------------
# UTILS
# ---------------------------------------------------------------------
def mkdir(p: Path):
    p.mkdir(parents=True, exist_ok=True)


def run(cmd):
    print(">>", " ".join(str(c) for c in cmd))
    subprocess.check_call(cmd)


def is_up_to_date(src, obj, dep):
    if "meta" in [src.stem, obj.stem, dep.stem]:
        return False
    if not obj.exists():
        return False
    if dep.exists() and dep.stat().st_mtime > obj.stat().st_mtime:
        return False
    return src.stat().st_mtime < obj.stat().st_mtime


# ---------------------------------------------------------------------
# BUILD RULES
# ---------------------------------------------------------------------
def compile_source(src: Path, is_test=False, target_platform=None):
    mkdir(OBJ_DIR)
    obj = OBJ_DIR / (src.stem + ".o")
    dep = OBJ_DIR / (src.stem + ".d")

    if is_up_to_date(src, obj, dep):
        return obj

    compiler, kind = detect_compiler(target_platform)
    
    if kind == "msvc":
        flags = MSVC_TEST_FLAGS if is_test else MSVC_FLAGS
        cmd = [compiler] + flags + ["/Fo" + str(obj), "/c", str(src),
                               f"/DGIT_COMMIT_HASH=\"{GIT_HASH}\""]
    else:
        if target_platform == "macos":
            flags = MACOS_FLAGS if not is_test else TEST_FLAGS + ["-mmacosx-version-min=10.13"]
        else:
            flags = TEST_FLAGS if is_test else COMMON_FLAGS
            
        cmd = [compiler] + flags + [
            f"-DGIT_COMMIT_HASH=\"{GIT_HASH}\"",
            "-MMD", "-MP",
            "-c", str(src), "-o", str(obj)
        ]

    run(cmd)
    return obj


def link_executable(objects, output: Path, target_platform=None):
    mkdir(BIN_DIR)
    compiler, kind = detect_compiler(target_platform)
    
    if kind == "msvc":
        cmd = [compiler] + list(map(str, objects)) + ["/Fe" + str(output)]
    else:
        cmd = [compiler] + list(map(str, objects)) + ["-lm", "-o", str(output)]
    run(cmd)


# ---------------------------------------------------------------------
# TEST CASE GENERATION
# ---------------------------------------------------------------------
def generate_tests():
    script = Path("../SLS_Tests/yaml_to_c_tests.py")
    yaml = Path("../SLS_Tests/cases.yaml")
    out = TEST_DIR / "lexer_tests.c"

    if not script.exists() or not yaml.exists():
        print("Test generation skipped (missing files).")
        return False

    if out.exists() and out.stat().st_mtime > yaml.stat().st_mtime:
        return True

    run([PYTHON, str(script), str(yaml), str(out)])
    return True


# ---------------------------------------------------------------------
# RP2040 BUILD
# ---------------------------------------------------------------------
def check_pico_sdk():
    """Check if Pico SDK is available"""
    sdk_path = Path(PICO_SDK_PATH)
    if not sdk_path.exists():
        print(f"ERROR: Pico SDK not found at {sdk_path}")
        print("Please set PICO_SDK_PATH environment variable or install SDK at ~/pico/pico-sdk")
        return False
    
    # Check for ARM toolchain
    if not shutil.which("arm-none-eabi-gcc"):
        print("ERROR: ARM toolchain not found!")
        print("Please install arm-none-eabi-gcc:")
        print("  Ubuntu/Debian: sudo apt install gcc-arm-none-eabi")
        print("  macOS: brew install arm-none-eabi-gcc")
        return False
    
    return True


def build_rp2040():
    if not check_pico_sdk():
        return False

    print("\n=== Building for RP2040 ===\n")

    print(f"PICO_SDK_PATH: {PICO_SDK_PATH}")

    generate_pico_headers()
    generate_pico_linker_scripts()

    mkdir(OBJ_DIR)
    mkdir(BIN_DIR)

    # Your application sources
    app_sources = [
        s for s in SRC_DIR.glob("*.c")
        if s.name not in ["main.c", "repl.c", "file.c"]
        and "test" not in s.stem.lower()
    ]

    pico_main = SRC_DIR / "pico_main.c"
    if not pico_main.exists():
        print("ERROR: src/pico_main.c is required for Pico builds")
        return False

    app_sources.append(pico_main)

    # SDK sources
    sdk_asm_sources = pico_sdk_asm_sources()
    sdk_sources = pico_sdk_sources()

    # Compile everything
    objects = []
    for src in sdk_asm_sources:
        objects.append(compile_pico_asm_source(src))
    for src in set(app_sources + sdk_sources):
        objects.append(compile_pico_source(src))

    # Link ELF
    elf = BIN_DIR / "sls_pico.elf"
    link_pico_elf(objects, elf)

    # Generate UF2
    uf2 = elf_to_uf2(elf)

    print("\nBuild successful!")
    print(f"ELF: {elf}")
    print(f"UF2: {uf2}")
    print("Flash by copying UF2 to Pico BOOTSEL drive")

    return True


# ---------------------------------------------------------------------
# MACOS BUILD
# ---------------------------------------------------------------------
def build_macos():
    """Build for macOS (can be run on macOS or cross-compile on Linux)"""
    print("\n=== Building for macOS ===\n")
    
    if PLATFORM != "macos" and PLATFORM != "linux":
        print("ERROR: macOS builds require macOS or Linux with osxcross")
        return False
    
    sources = list(SRC_DIR.glob("*.c"))
    objects = [compile_source(s, is_test=False, target_platform="macos") for s in sources]
    
    macos_target = BIN_DIR / "sls_macos"
    link_executable(objects, macos_target, target_platform="macos")
    
    print(f"\nBuild successful! Binary: {macos_target}")
    return True


# ---------------------------------------------------------------------
# TARGETS
# ---------------------------------------------------------------------
def build_main():
    sources = list(SRC_DIR.glob("*.c"))
    objects = [compile_source(s, is_test=False) for s in sources if s.name != 'pico_main.c']
    link_executable(objects, TARGET)


def build_tests():
    generate_tests()
    test_sources = list(TEST_DIR.glob("*.c"))
    main_sources = [s for s in SRC_DIR.glob("*.c") if s.name != "main.c"]

    test_objects = [compile_source(s, is_test=True) for s in test_sources]
    shared_objects = [compile_source(s, is_test=False) for s in main_sources]

    link_executable(test_objects + shared_objects, TEST_TARGET)


def clean():
    shutil.rmtree(OBJ_DIR, ignore_errors=True)
    shutil.rmtree(BIN_DIR, ignore_errors=True)
    shutil.rmtree(PICO_GENERATED, ignore_errors=True)
    print("Cleaned.")


def run_main():
    build_main()
    print("\n--- Running program ---\n")
    subprocess.call([str(TARGET)])


def run_tests():
    build_tests()
    print("\n--- Running tests ---\n")
    subprocess.call([str(TEST_TARGET)])


def debug_tests():
    build_tests()
    if os.name == "nt":
        print("Debugging on Windows requires Visual Studio debugger.")
    else:
        subprocess.call(["gdb", str(TEST_TARGET)])


def show_help():
    help_text = """
Build Script for Multi-Platform Compilation
============================================

Usage: python3 build.py [command]

Commands:
  all, main, build    Build for current platform
  run                 Build and run program
  test                Build and run tests
  debug               Build tests and run debugger
  clean               Remove build artifacts
  
  macos               Build for macOS
  pico, rp2040        Build for RP2040 (Raspberry Pi Pico)
  
  help                Show this help message

Environment Variables:
  PICO_SDK_PATH       Path to Pico SDK (default: ~/pico/pico-sdk)

Examples:
  python3 build.py build      # Build for current platform
  python3 build.py pico       # Build for RP2040
  python3 build.py macos      # Build for macOS
  python3 build.py clean      # Clean all build files
"""
    print(help_text)


# ---------------------------------------------------------------------
# ENTRY POINT
# ---------------------------------------------------------------------
def main():
    if len(sys.argv) < 2:
        show_help()
        return

    cmd = sys.argv[1]

    match cmd:
        case "all" | "main" | "build":
            build_main()
        case "run":
            run_main()
        case "test":
            run_tests()
        case "debug":
            debug_tests()
        case "clean":
            clean()
        case "macos":
            build_macos()
        case "pico" | "rp2040":
            build_rp2040()
        case "help" | "-h" | "--help":
            show_help()
        case _:
            print(f"Unknown target: {cmd}")
            print("Run 'python3 build.py help' for usage information")


if __name__ == "__main__":
    main()