#!/usr/bin/env python3
"""Convert an ASCII AIGER (.aag) file to Graphviz DOT format.

Usage: python3 scripts/aag2dot.py circuits/spec.aag | dot -Tsvg -o diagrams/spec.svg
"""

import sys
from pathlib import Path


def parse_aag(text: str) -> dict:
    lines = text.strip().split('\n')
    header = lines[0].split()
    assert header[0] == 'aag'
    max_var = int(header[1])
    n_inputs = int(header[2])
    n_latches = int(header[3])
    n_outputs = int(header[4])
    n_ands = int(header[5])

    idx = 1
    inputs = []
    for i in range(n_inputs):
        inputs.append(int(lines[idx]))
        idx += 1

    latches = []
    for i in range(n_latches):
        parts = lines[idx].split()
        latches.append((int(parts[0]), int(parts[1])))
        idx += 1

    outputs = []
    for i in range(n_outputs):
        outputs.append(int(lines[idx]))
        idx += 1

    ands = []
    for i in range(n_ands):
        parts = lines[idx].split()
        ands.append((int(parts[0]), int(parts[1]), int(parts[2])))
        idx += 1

    # Symbol table
    input_names = {}
    output_names = {}
    latch_names = {}
    while idx < len(lines):
        line = lines[idx]
        if line.startswith('i'):
            parts = line.split(' ', 1)
            input_names[int(parts[0][1:])] = parts[1]
        elif line.startswith('o'):
            parts = line.split(' ', 1)
            output_names[int(parts[0][1:])] = parts[1]
        elif line.startswith('l'):
            parts = line.split(' ', 1)
            latch_names[int(parts[0][1:])] = parts[1]
        elif line.startswith('c'):
            break
        idx += 1

    return {
        'inputs': inputs,
        'latches': latches,
        'outputs': outputs,
        'ands': ands,
        'input_names': input_names,
        'output_names': output_names,
        'latch_names': latch_names,
    }


def lit_node(lit: int) -> tuple[str, bool]:
    """Return (node_id, is_negated) for a literal."""
    var = lit >> 1
    negated = lit & 1
    return f'n{var}', bool(negated)


def to_dot(aag: dict) -> str:
    lines = ['digraph aiger {']
    lines.append('  rankdir=LR;')
    lines.append('  node [shape=record];')
    lines.append('')

    # Constant 0
    lines.append('  n0 [label="0", shape=circle, style=filled, fillcolor=gray90];')

    # Inputs
    lines.append('  { rank=source;')
    for i, inp in enumerate(aag['inputs']):
        var = inp >> 1
        name = aag['input_names'].get(i, f'i{i}')
        lines.append(f'    n{var} [label="{name}", shape=box, style=filled, fillcolor=lightblue];')
    lines.append('  }')

    # Latches
    for i, (latch_lit, next_lit) in enumerate(aag['latches']):
        var = latch_lit >> 1
        name = aag['latch_names'].get(i, f'L{i}')
        lines.append(f'  n{var} [label="{name}", shape=box, style="filled,rounded", fillcolor=lightyellow];')

    # AND gates
    for lhs, rhs0, rhs1 in aag['ands']:
        var = lhs >> 1
        lines.append(f'  n{var} [label="&", shape=circle, style=filled, fillcolor=white];')

    # Outputs
    lines.append('  { rank=sink;')
    for i, out_lit in enumerate(aag['outputs']):
        name = aag['output_names'].get(i, f'o{i}')
        lines.append(f'    out_{i} [label="{name}", shape=box, style=filled, fillcolor=lightgreen];')
    lines.append('  }')

    lines.append('')

    # AND gate edges
    for lhs, rhs0, rhs1 in aag['ands']:
        lhs_var = lhs >> 1
        for rhs in [rhs0, rhs1]:
            src_node, negated = lit_node(rhs)
            style = 'style=dashed, color=red' if negated else ''
            label = 'label="~"' if negated else ''
            attrs = ', '.join(filter(None, [style, label]))
            if attrs:
                lines.append(f'  {src_node} -> n{lhs_var} [{attrs}];')
            else:
                lines.append(f'  {src_node} -> n{lhs_var};')

    # Latch feedback edges
    for i, (latch_lit, next_lit) in enumerate(aag['latches']):
        latch_var = latch_lit >> 1
        src_node, negated = lit_node(next_lit)
        style = 'style=dashed, color=red' if negated else ''
        label = 'label="~"' if negated else ''
        extra = 'constraint=false, '
        attrs = ', '.join(filter(None, [extra + style, label]))
        lines.append(f'  {src_node} -> n{latch_var} [{attrs}];')

    # Output edges
    for i, out_lit in enumerate(aag['outputs']):
        src_node, negated = lit_node(out_lit)
        style = 'style=dashed, color=red' if negated else ''
        label = 'label="~"' if negated else ''
        attrs = ', '.join(filter(None, [style, label]))
        if attrs:
            lines.append(f'  {src_node} -> out_{i} [{attrs}];')
        else:
            lines.append(f'  {src_node} -> out_{i};')

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


def main():
    if len(sys.argv) < 2:
        print(f'Usage: {sys.argv[0]} <file.aag>', file=sys.stderr)
        sys.exit(1)

    text = Path(sys.argv[1]).read_text()
    aag = parse_aag(text)
    print(to_dot(aag))


if __name__ == '__main__':
    main()