"""
C code generator for verified arena allocator specs.

Takes a verified ArenaSpec and produces a complete, compilable C
implementation. The generated code is structured to match the
properties that were formally verified.
"""

from .spec import ArenaSpec, AllocStrategy, GrowthPolicy, ThreadSafety


def generate_c(spec: ArenaSpec) -> str:
    """Generate complete C implementation from a verified ArenaSpec."""
    parts = [
        _generate_header(spec),
        _generate_struct(spec),
        _generate_create(spec),
        _generate_alloc(spec),
        _generate_reset(spec),
        _generate_destroy(spec),
    ]
    return "\n".join(parts)


def _generate_header(spec: ArenaSpec) -> str:
    name = spec.name.upper()
    includes = [
        "#include <stddef.h>",
        "#include <stdint.h>",
        "#include <stdlib.h>",
    ]
    if spec.safety.zero_on_alloc or spec.safety.zero_on_reset:
        includes.append("#include <string.h>")
    if spec.thread_safety == ThreadSafety.LOCK_FREE:
        includes.append("#include <stdatomic.h>")

    return f"""/*
 * Auto-generated arena allocator: {spec.name}
 *
 * Strategy:    {spec.strategy.value}
 * Page size:   {spec.memory.page_size} bytes
 * Max pages:   {spec.memory.max_pages}
 * Alignment:   {spec.alignment.min_align} bytes
 * Threading:   {spec.thread_safety.value}
 *
 * VERIFIED PROPERTIES:
 *   - Bump pointer cannot overflow when space available
 *   - All allocations are {spec.alignment.min_align}-byte aligned
 *   - Total memory bounded at {spec.memory.total_capacity} bytes
 *
 * Generated by prompt2bin — do not edit manually.
 */

#ifndef {name}_H
#define {name}_H

{chr(10).join(includes)}

""".lstrip()


def _generate_struct(spec: ArenaSpec) -> str:
    offset_type = "atomic_size_t" if spec.thread_safety == ThreadSafety.LOCK_FREE else "size_t"

    struct = f"""typedef struct {{
    uint8_t *buf;       /* backing memory */
    {offset_type} offset;  /* current bump pointer offset */
    size_t capacity;    /* usable capacity (page_size - header) */"""

    if spec.safety.no_use_after_reset:
        struct += "\n    uint64_t generation; /* generation counter for use-after-reset detection */"

    if spec.growth != GrowthPolicy.FIXED:
        struct += f"""
    struct {spec.name}_t *next; /* linked list for chained arenas */
    size_t page_count;          /* number of pages allocated */"""

    struct += f"""
}} {spec.name}_t;

"""
    return struct


def _generate_create(spec: ArenaSpec) -> str:
    page_size = spec.memory.page_size
    header_size = spec.memory.header_size
    align = spec.alignment.min_align
    name = spec.name

    # Use aligned_alloc for guaranteed alignment of the arena base
    return f"""{name}_t *{name}_create(void) {{
    /* Allocate arena struct + backing buffer in a single allocation
     * to minimize fragmentation and ensure locality. */
    size_t total = sizeof({name}_t) + {page_size};

    /* aligned_alloc requires size to be a multiple of alignment */
    size_t alloc_size = (total + {align - 1}) & ~(size_t){align - 1};
    void *mem = aligned_alloc({align}, alloc_size);
    if (!mem) return NULL;

    {name}_t *arena = ({name}_t *)mem;
    arena->buf = (uint8_t *)mem + sizeof({name}_t);

    /* Ensure buf is aligned — adjust if struct size isn't a multiple of alignment */
    uintptr_t buf_addr = (uintptr_t)arena->buf;
    buf_addr = (buf_addr + {align - 1}) & ~(uintptr_t){align - 1};
    arena->buf = (uint8_t *)buf_addr;

    arena->offset = 0;
    arena->capacity = {page_size - header_size};
{"    arena->generation = 0;" if spec.safety.no_use_after_reset else ""}
    return arena;
}}

"""


def _generate_alloc(spec: ArenaSpec) -> str:
    name = spec.name
    align = spec.alignment.min_align
    align_mask = align - 1

    if spec.thread_safety == ThreadSafety.LOCK_FREE:
        return _generate_alloc_lock_free(spec)

    zero_line = ""
    if spec.safety.zero_on_alloc:
        zero_line = "    memset(ptr, 0, size);\n"

    return f"""void *{name}_alloc({name}_t *arena, size_t size) {{
    if (size == 0 || size > {spec.memory.effective_max_alloc}) return NULL;

    /* Round offset up to {align}-byte alignment — verified by Z3:
     * for any aligned base and any offset, (offset + mask) & ~mask
     * produces an aligned result. */
    size_t aligned_offset = (arena->offset + {align_mask}) & ~(size_t){align_mask};

    /* Bounds check — verified by Z3: if this check passes,
     * aligned_offset + size cannot overflow. */
    if (aligned_offset + size > arena->capacity) return NULL;

    void *ptr = arena->buf + aligned_offset;
    arena->offset = aligned_offset + size;
{zero_line}    return ptr;
}}

"""


def _generate_alloc_lock_free(spec: ArenaSpec) -> str:
    name = spec.name
    align = spec.alignment.min_align
    align_mask = align - 1

    return f"""void *{name}_alloc({name}_t *arena, size_t size) {{
    if (size == 0 || size > {spec.memory.effective_max_alloc}) return NULL;

    size_t current, aligned, new_offset;
    do {{
        current = atomic_load_explicit(&arena->offset, memory_order_relaxed);
        aligned = (current + {align_mask}) & ~(size_t){align_mask};
        new_offset = aligned + size;
        if (new_offset > arena->capacity) return NULL;
    }} while (!atomic_compare_exchange_weak_explicit(
        &arena->offset, &current, new_offset,
        memory_order_release, memory_order_relaxed));

    return arena->buf + aligned;
}}

"""


def _generate_reset(spec: ArenaSpec) -> str:
    name = spec.name
    lines = [f"void {name}_reset({name}_t *arena) {{"]

    if spec.safety.zero_on_reset:
        lines.append("    memset(arena->buf, 0, arena->offset);")

    if spec.thread_safety == ThreadSafety.LOCK_FREE:
        lines.append("    atomic_store_explicit(&arena->offset, 0, memory_order_release);")
    else:
        lines.append("    arena->offset = 0;")

    if spec.safety.no_use_after_reset:
        lines.append("    arena->generation++;")

    lines.append("}")
    lines.append("")
    return "\n".join(lines)


def _generate_destroy(spec: ArenaSpec) -> str:
    name = spec.name
    footer = f"""void {name}_destroy({name}_t *arena) {{
    if (arena) free(arena);
}}

#endif /* {name.upper()}_H */
"""
    return footer
