#!/usr/bin/env python3
"""
CCT-Lang Compiler (cctc) – translates CCT-Lang to C.
Refactored to support blocks, type aliases, struct-arrays, nested for-loops,
ternary operators (? :), pointer types (T*), dereferencing (*expr), and address-of (&expr).
"""

import sys
import re
import argparse
from collections import namedtuple

# Tokens
TOKENS = [
    ('NUMBER', r'\d+(?:\.\d+)?'),
    ('IDENT', r'[a-zA-Z_][a-zA-Z0-9_]*'),
    ('STRING', r'"[^"]*"'),
    ('LPAREN', r'\('),
    ('RPAREN', r'\)'),
    ('LBRACE', r'\{'),
    ('RBRACE', r'\}'),
    ('LBRACK', r'\['),
    ('RBRACK', r'\]'),
    ('SEMICOLON', r';'),
    ('COMMA', r','),
    ('COMMENT', r'//[^\n]*'),
    ('DOTDOT', r'\.\.'),
    ('OPERATOR', r'==|!=|<=|>=|&&|\|\||\+\+|--|\+=|-=|\*=|/=|%=|[=+*/<>!%&?-]'),
    ('COLON', r':'),
    ('WHITESPACE', r'\s+'),
]

KEYWORDS = {
    'type', 'array', 'function', 'return', 'if', 'else', 'for', 'while', 'in', 'then',
    'prob_decimal', 'prob_float', 'superbool', 'tseries', 'lazy', 'uint8',
    'uint32', 'float', 'int', 'bool', 'true', 'false', 'print', 'void',
    'observe', 'measure', 'collapse_if', 'uniform', 'gate'
}

class Token:
    def __init__(self, type, value, line, col):
        self.type = type
        self.value = value
        self.line = line
        self.col = col
    def __repr__(self):
        return f'Token({self.type}, {self.value})'

def lex(code):
    tokens = []
    pos = 0
    line = 1
    col = 1
    while pos < len(code):
        match = None
        for tok_type, pattern in TOKENS:
            regex = re.compile(pattern)
            m = regex.match(code, pos)
            if m:
                match = m
                value = m.group(0)
                if tok_type == 'COMMENT':
                    col += len(value)
                elif tok_type == 'WHITESPACE':
                    lines = value.count('\n')
                    if lines:
                        line += lines
                        col = len(value) - value.rfind('\n')
                    else:
                        col += len(value)
                else:
                    if tok_type == 'IDENT' and value in KEYWORDS:
                        tok_type = 'KEYWORD'
                    tokens.append(Token(tok_type, value, line, col))
                    col += len(value)
                pos = m.end()
                break
        if not match:
            raise SyntaxError(f'Unexpected character {code[pos]} at {line}:{col}')
    return tokens

# AST Nodes
class ASTNode: pass
class Program(ASTNode):
    def __init__(self, declarations): self.declarations = declarations
class Function(ASTNode):
    def __init__(self, ret_type, name, params, body):
        self.ret_type = ret_type; self.name = name; self.params = params; self.body = body
class VarDecl(ASTNode):
    def __init__(self, type_spec, name, init=None):
        self.type_spec = type_spec; self.name = name; self.init = init
class TypeDef(ASTNode):
    def __init__(self, name, type_spec):
        self.name = name; self.type_spec = type_spec
class Block(ASTNode):
    def __init__(self, stmts): self.stmts = stmts
class IfStmt(ASTNode):
    def __init__(self, cond, then_stmt, else_stmt=None):
        self.cond = cond; self.then_stmt = then_stmt; self.else_stmt = else_stmt
class ForRange(ASTNode):
    def __init__(self, var, start, end):
        self.var = var; self.start = start; self.end = end
class ForStmt(ASTNode):
    def __init__(self, ranges, body):
        self.ranges = ranges; self.body = body
class WhileStmt(ASTNode):
    def __init__(self, cond, body): self.cond = cond; self.body = body
class ReturnStmt(ASTNode):
    def __init__(self, expr): self.expr = expr
class PrintStmt(ASTNode):
    def __init__(self, args): self.args = args
class Assignment(ASTNode):
    def __init__(self, target, value): self.target = target; self.value = value
class Expr(ASTNode): pass
class BinOp(Expr):
    def __init__(self, left, op, right): self.left = left; self.op = op; self.right = right
class UnaryOp(Expr):
    def __init__(self, op, expr, postfix=False): self.op = op; self.expr = expr; self.postfix = postfix
class TernaryExpr(Expr):
    def __init__(self, cond, then_expr, else_expr):
        self.cond = cond; self.then_expr = then_expr; self.else_expr = else_expr
class Call(Expr):
    def __init__(self, func, args): self.func = func; self.args = args
class Var(Expr):
    def __init__(self, name): self.name = name
class ArrayIndex(Expr):
    def __init__(self, target, index): self.target = target; self.index = index
class Literal(Expr):
    def __init__(self, value, kind): self.value = value; self.kind = kind

class TypeSpec(ASTNode): pass
class BasicType(TypeSpec):
    def __init__(self, name): self.name = name
class PointerType(TypeSpec):
    def __init__(self, base_type): self.base_type = base_type
class StructArrayType(TypeSpec):
    def __init__(self, base_type, size): self.base_type = base_type; self.size = size
class ArrayType(TypeSpec):
    def __init__(self, base_type, size): self.base_type = base_type; self.size = size
class ProbType(TypeSpec):
    def __init__(self, p, s): self.p = p; self.s = s

# Parser
class Parser:
    def __init__(self, tokens):
        self.tokens = tokens
        self.pos = 0

    def peek(self, n=0):
        if self.pos + n < len(self.tokens): return self.tokens[self.pos + n]
        return None

    def consume(self, expected_type=None, expected_value=None):
        tok = self.peek()
        if not tok: raise SyntaxError('Unexpected end of input')
        if expected_type and tok.type != expected_type:
            raise SyntaxError(f'Expected {expected_type}, got {tok.type} at {tok.line}:{tok.col}')
        if expected_value and tok.value != expected_value:
            raise SyntaxError(f'Expected {expected_value}, got {tok.value} at {tok.line}:{tok.col}')
        self.pos += 1
        return tok

    def parse(self):
        decls = []
        while self.peek():
            decls.append(self.parse_declaration())
        return Program(decls)

    def parse_declaration(self):
        tok = self.peek()
        if tok.type == 'KEYWORD' and tok.value == 'type':
            self.consume('KEYWORD', 'type')
            name = self.consume('IDENT').value
            self.consume('OPERATOR', '=')
            type_spec = self.parse_type()
            self.consume('SEMICOLON')
            return TypeDef(name, type_spec)
        elif tok.type == 'KEYWORD' and tok.value == 'function':
            return self.parse_function()
        else:
            # Check if it's a function: type_spec IDENT (
            type_spec = self.parse_type()
            name = self.consume('IDENT').value
            if self.peek() and self.peek().type == 'LPAREN':
                return self.finish_parse_function(type_spec, name)
            
            while self.peek() and self.peek().type == 'LBRACK':
                self.consume('LBRACK')
                size = self.parse_expression() if self.peek().type != 'RBRACK' else None
                self.consume('RBRACK')
                type_spec = ArrayType(type_spec, size)
            
            init = None
            if self.peek() and self.peek().value == '=':
                self.consume('OPERATOR', '=')
                if self.peek().type == 'LBRACE':
                    init = self.parse_initializer_list()
                else:
                    init = self.parse_expression()
            self.consume('SEMICOLON')
            return VarDecl(type_spec, name, init)

    def parse_initializer_list(self):
        self.consume('LBRACE')
        elems = []
        while self.peek() and self.peek().type != 'RBRACE':
            elems.append(self.parse_expression())
            if self.peek() and self.peek().type == 'COMMA':
                self.consume('COMMA')
        self.consume('RBRACE')
        return Literal(elems, 'INIT_LIST')

    def parse_type(self):
        tok = self.peek()
        if tok.type == 'KEYWORD' and tok.value == 'prob_decimal':
            self.consume()
            self.consume('LPAREN')
            p = self.consume('NUMBER').value
            self.consume('COMMA')
            s = self.consume('NUMBER').value
            self.consume('RPAREN')
            t = ProbType(p, s)
        elif tok.type == 'KEYWORD' and tok.value == 'array':
            self.consume()
            self.consume('OPERATOR', '<')
            base = self.parse_type()
            self.consume('COMMA')
            size = self.parse_expression(5)
            self.consume('OPERATOR', '>')
            t = StructArrayType(base, size)
        else:
            tok = self.consume()
            if tok.type not in ('IDENT', 'KEYWORD'):
                raise SyntaxError(f'Expected type name, got {tok.type} at {tok.line}:{tok.col}')
            name = tok.value
            t = BasicType(name)
        
        while True:
            tok = self.peek()
            if tok and tok.type == 'OPERATOR' and tok.value == '*':
                self.consume('OPERATOR', '*')
                t = PointerType(t)
            elif tok and tok.type == 'LBRACK':
                self.consume('LBRACK')
                size = self.parse_expression() if self.peek().type != 'RBRACK' else None
                self.consume('RBRACK')
                t = ArrayType(t, size)
            else:
                break
        return t

    def parse_function(self):
        self.consume('KEYWORD', 'function')
        # Heuristic: if next is IDENT and then LPAREN, return type is void
        if self.peek().type == 'IDENT' and self.peek(1) and self.peek(1).type == 'LPAREN':
            ret_type = BasicType('void')
            name = self.consume('IDENT').value
        else:
            ret_type = self.parse_type()
            name = self.consume('IDENT').value
        return self.finish_parse_function(ret_type, name)

    def finish_parse_function(self, ret_type, name):
        self.consume('LPAREN')
        params = []
        while self.peek() and self.peek().type != 'RPAREN':
            ptype = self.parse_type()
            pname = self.consume('IDENT').value
            while self.peek() and self.peek().type == 'LBRACK':
                self.consume('LBRACK')
                size = self.parse_expression() if self.peek().type != 'RBRACK' else None
                self.consume('RBRACK')
                ptype = ArrayType(ptype, size)
            params.append((ptype, pname))
            if self.peek() and self.peek().type == 'COMMA':
                self.consume('COMMA')
        self.consume('RPAREN')
        body = self.parse_statement()
        if not isinstance(body, Block):
            body = Block([body])
        return Function(ret_type, name, params, body)

    def parse_statement(self):
        tok = self.peek()
        if not tok: return None
        if tok.type == 'LBRACE':
            self.consume('LBRACE')
            stmts = []
            while self.peek() and self.peek().type != 'RBRACE':
                stmts.append(self.parse_statement())
            self.consume('RBRACE')
            return Block(stmts)
        elif tok.type == 'KEYWORD':
            if tok.value == 'if': return self.parse_if()
            if tok.value == 'for': return self.parse_for()
            if tok.value == 'while': return self.parse_while()
            if tok.value == 'return':
                self.consume()
                expr = self.parse_expression()
                self.consume('SEMICOLON')
                return ReturnStmt(expr)
            if tok.value == 'print':
                self.consume()
                self.consume('LPAREN')
                args = []
                while self.peek() and self.peek().type != 'RPAREN':
                    args.append(self.parse_expression())
                    if self.peek() and self.peek().type == 'COMMA': self.consume('COMMA')
                self.consume('RPAREN')
                self.consume('SEMICOLON')
                return PrintStmt(args)
            if tok.value in {'uint8', 'uint32', 'float', 'int', 'bool', 'type', 'prob_decimal', 'array'}:
                return self.parse_declaration()
        
        if tok.type == 'IDENT' and self.peek(1) and self.peek(1).type == 'IDENT':
            return self.parse_declaration()
        
        expr = self.parse_expression()
        if self.peek() and self.peek().value == '=':
            self.consume('OPERATOR', '=')
            if self.peek().type == 'LBRACE':
                value = self.parse_initializer_list()
            else:
                value = self.parse_expression()
            self.consume('SEMICOLON')
            return Assignment(expr, value)
        self.consume('SEMICOLON')
        return expr

    def parse_if(self):
        self.consume('KEYWORD', 'if')
        if self.peek().type == 'LPAREN':
            self.consume('LPAREN')
            cond = self.parse_expression()
            self.consume('RPAREN')
        else:
            cond = self.parse_expression()
        if self.peek() and self.peek().value == 'then':
            self.consume('KEYWORD', 'then')
        then_stmt = self.parse_statement()
        else_stmt = None
        if self.peek() and self.peek().value == 'else':
            self.consume('KEYWORD', 'else')
            else_stmt = self.parse_statement()
        return IfStmt(cond, then_stmt, else_stmt)

    def parse_for(self):
        self.consume('KEYWORD', 'for')
        ranges = []
        while True:
            var = self.consume('IDENT').value
            self.consume('KEYWORD', 'in')
            start = self.parse_expression()
            self.consume('DOTDOT')
            end = self.parse_expression()
            ranges.append(ForRange(var, start, end))
            if self.peek() and self.peek().type == 'COMMA':
                self.consume('COMMA')
            else:
                break
        body = self.parse_statement()
        return ForStmt(ranges, body)

    def parse_while(self):
        self.consume('KEYWORD', 'while')
        self.consume('LPAREN')
        cond = self.parse_expression()
        self.consume('RPAREN')
        body = self.parse_statement()
        return WhileStmt(cond, body)

    def parse_expression(self, min_prec=0):
        left = self.parse_unary()
        while True:
            tok = self.peek()
            if not tok: break
            
            # Handle Ternary (cond ? then : else)
            if tok.type == 'OPERATOR' and tok.value == '?':
                if min_prec > 0: break 
                self.consume('OPERATOR', '?')
                then_expr = self.parse_expression(0)
                self.consume('COLON')
                else_expr = self.parse_expression(0)
                left = TernaryExpr(left, then_expr, else_expr)
                continue

            if tok.type != 'OPERATOR': break
            op = tok.value
            prec = self.precedence(op)
            if prec < min_prec: break
            self.consume()
            right = self.parse_expression(prec + 1)
            left = BinOp(left, op, right)
        return left

    def precedence(self, op):
        if op in ('=', '+=', '-=', '*=', '/=', '%='): return 0
        if op in ('||',): return 1
        if op in ('&&',): return 2
        if op in ('==', '!='): return 3
        if op in ('<', '>', '<=', '>='): return 4
        if op in ('+', '-'): return 5
        if op in ('*', '/', '%'): return 6
        return 0

    def parse_unary(self):
        tok = self.peek()
        if tok and tok.type == 'OPERATOR' and tok.value in ('*', '&', '!', '-', '+', '++', '--'):
            self.consume()
            return UnaryOp(tok.value, self.parse_unary(), postfix=False)
        return self.parse_primary()

    def parse_primary(self):
        tok = self.peek()
        if tok.type == 'NUMBER':
            self.consume()
            return Literal(tok.value, 'NUMBER')
        if tok.type == 'STRING':
            self.consume()
            return Literal(tok.value[1:-1], 'STRING')
        if tok.type == 'KEYWORD' and tok.value in ('true', 'false'):
            self.consume()
            return Literal(tok.value == 'true', 'BOOL')
        if tok.type == 'IDENT':
            name = tok.value
            self.consume()
            node = Var(name)
            while self.peek():
                if self.peek().type == 'LPAREN':
                    self.consume('LPAREN')
                    args = []
                    while self.peek() and self.peek().type != 'RPAREN':
                        args.append(self.parse_expression())
                        if self.peek() and self.peek().type == 'COMMA': self.consume('COMMA')
                    self.consume('RPAREN')
                    node = Call(name, args)
                elif self.peek().type == 'LBRACK':
                    self.consume('LBRACK')
                    idx = self.parse_expression()
                    self.consume('RBRACK')
                    node = ArrayIndex(node, idx)
                elif self.peek().type == 'OPERATOR' and self.peek().value in ('++', '--'):
                    op = self.consume().value
                    node = UnaryOp(op, node, postfix=True)
                else:
                    break
            return node
        if tok.value == '(':
            self.consume('LPAREN')
            expr = self.parse_expression()
            self.consume('RPAREN')
            return expr
        raise SyntaxError(f'Unexpected token {tok} at {tok.line}:{tok.col}')

# Scope & Generation
class Scope:
    def __init__(self, parent=None):
        self.parent = parent
        self.vars = {}
    def define(self, name, type_spec): self.vars[name] = type_spec
    def lookup(self, name):
        if name in self.vars: return self.vars[name]
        if self.parent: return self.parent.lookup(name)
        return None

class CGenerator:
    def __init__(self):
        self.code = []
        self.indent = 0
        self.typedefs = {}
        self.global_scope = Scope()
        self.current_scope = self.global_scope

    def emit(self, line): self.code.append('    ' * self.indent + line)

    def get_type(self, node):
        if isinstance(node, Var):
            return self.current_scope.lookup(node.name)
        if isinstance(node, ArrayIndex):
            target_type = self.get_type(node.target)
            while isinstance(target_type, BasicType) and target_type.name in self.typedefs:
                target_type = self.typedefs[target_type.name]
            if isinstance(target_type, (StructArrayType, ArrayType, PointerType)):
                return target_type.base_type
        if isinstance(node, UnaryOp) and node.op == '*':
            target_type = self.get_type(node.expr)
            while isinstance(target_type, BasicType) and target_type.name in self.typedefs:
                target_type = self.typedefs[target_type.name]
            if isinstance(target_type, PointerType):
                return target_type.base_type
        if isinstance(node, UnaryOp) and node.op == '&':
            target_type = self.get_type(node.expr)
            if target_type:
                return PointerType(target_type)
        return None

    def generate(self, node):
        if isinstance(node, Program):
            self.emit('#include "cct_runtime.h"')
            self.emit('#include <stdio.h>')
            self.emit('#include <stdint.h>')
            self.emit('#include <stdbool.h>')
            self.emit('#include <math.h>')
            self.emit('')
            for decl in node.declarations:
                if decl: self.generate(decl)
        elif isinstance(node, TypeDef):
            self.typedefs[node.name] = node.type_spec
            if isinstance(node.type_spec, StructArrayType):
                bt = self.c_type(node.type_spec.base_type)
                sz = self.expr_str(node.type_spec.size)
                self.emit(f'typedef struct {{ {bt} data[{sz}]; }} {node.name};')
            else:
                self.emit(f'typedef {self.c_type(node.type_spec)} {node.name};')
        elif isinstance(node, Function):
            self.current_scope = Scope(self.global_scope)
            ret_c = self.c_type(node.ret_type)
            # Rename functions that conflict with C standard library
            name = node.name
            if name in ('exp', 'log', 'log10', 'pow', 'sqrt'):
                name = f'cct_user_{name}'
            
            params_c = [self.c_decl(pt, pn) for pt, pn in node.params]
            for pt, pn in node.params: self.current_scope.define(pn, pt)
            self.emit(f'{ret_c} {name}({", ".join(params_c)}) {{')
            self.indent += 1
            self.generate(node.body)
            self.indent -= 1
            self.emit('}')
            self.emit('')
            self.current_scope = self.global_scope
        elif isinstance(node, Block):
            for stmt in node.stmts: self.generate(stmt)
        elif isinstance(node, VarDecl):
            self.current_scope.define(node.name, node.type_spec)
            init_c = f' = {self.expr_str(node.init)}' if node.init else ''
            self.emit(f'{self.c_decl(node.type_spec, node.name)}{init_c};')
        elif isinstance(node, Assignment):
            self.emit(f'{self.expr_str(node.target)} = {self.expr_str(node.value)};')
        elif isinstance(node, IfStmt):
            self.emit(f'if ({self.expr_str(node.cond)}) {{')
            self.indent += 1
            self.generate(node.then_stmt)
            self.indent -= 1
            if node.else_stmt:
                self.emit('} else {')
                self.indent += 1
                self.generate(node.else_stmt)
                self.indent -= 1
            self.emit('}')
        elif isinstance(node, ForStmt):
            for r in node.ranges:
                self.emit(f'for (int {r.var} = {self.expr_str(r.start)}; {r.var} <= {self.expr_str(r.end)}; {r.var}++) {{')
                self.indent += 1
            self.generate(node.body)
            for _ in node.ranges:
                self.indent -= 1
                self.emit('}')
        elif isinstance(node, WhileStmt):
            self.emit(f'while ({self.expr_str(node.cond)}) {{')
            self.indent += 1
            self.generate(node.body)
            self.indent -= 1
            self.emit('}')
        elif isinstance(node, ReturnStmt):
            self.emit(f'return {self.expr_str(node.expr)};')
        elif isinstance(node, PrintStmt):
            fmt, args = "", []
            for arg in node.args:
                if isinstance(arg, Literal) and arg.kind == 'STRING':
                    fmt += arg.value.replace('%', '%%')
                else:
                    fmt += "%g"
                    args.append(f"(double)({self.expr_str(arg)})")
            self.emit(f'cct_print("{fmt}"{", " + ", ".join(args) if args else ""});')
        elif isinstance(node, Expr):
            self.emit(f'{self.expr_str(node)};')

    def expr_str(self, node):
        if isinstance(node, BinOp):
            return f'({self.expr_str(node.left)} {node.op} {self.expr_str(node.right)})'
        if isinstance(node, UnaryOp):
            if node.postfix:
                return f'({self.expr_str(node.expr)}{node.op})'
            return f'({node.op}{self.expr_str(node.expr)})'
        if isinstance(node, TernaryExpr):
            return f'({self.expr_str(node.cond)} ? {self.expr_str(node.then_expr)} : {self.expr_str(node.else_expr)})'
        if isinstance(node, Call):
            args = [self.expr_str(a) for a in node.args]
            f_name = node.func
            if f_name in ('addp', 'multp', 'logp', 'exp', 'start_timer', 'stop_timer', 'builtin_log', 'builtin_exp'):
                real_name = f_name[8:] if f_name.startswith('builtin_') else f_name
                return f'cct_{real_name}({", ".join(args)})'
            return f'{f_name}({", ".join(args)})'
        if isinstance(node, Var): return node.name
        if isinstance(node, ArrayIndex):
            t = self.get_type(node.target)
            while isinstance(t, BasicType) and t.name in self.typedefs: t = self.typedefs[t.name]
            if isinstance(t, StructArrayType):
                return f'{self.expr_str(node.target)}.data[{self.expr_str(node.index)}]'
            return f'{self.expr_str(node.target)}[{self.expr_str(node.index)}]'
        if isinstance(node, Literal):
            if node.kind == 'NUMBER': return str(node.value)
            if node.kind == 'STRING': return f'"{node.value}"'
            if node.kind == 'BOOL': return 'true' if node.value else 'false'
            if node.kind == 'INIT_LIST': return '{' + ', '.join(self.expr_str(e) for e in node.value) + '}'
        return ''

    def c_type(self, type_spec):
        if isinstance(type_spec, BasicType):
            m = {'uint8': 'uint8_t', 'uint32': 'uint32_t', 'float': 'float', 'int': 'int', 'bool': 'bool', 'void': 'void', 'timer': 'uint32_t'}
            return m.get(type_spec.name, type_spec.name)
        if isinstance(type_spec, ProbType): return 'cct_prob_t'
        if isinstance(type_spec, ArrayType): return self.c_type(type_spec.base_type)
        if isinstance(type_spec, PointerType): return f'{self.c_type(type_spec.base_type)}*'
        return 'void'

    def c_decl(self, type_spec, name):
        if isinstance(type_spec, ArrayType):
            base, dims = type_spec, []
            while isinstance(base, ArrayType):
                dims.append(self.expr_str(base.size) if base.size else "")
                base = base.base_type
            # Reverse dims to match C array order (outermost dimension first)
            return f"{self.c_type(base)} {name}{''.join(f'[{d}]' for d in reversed(dims))}"
        if isinstance(type_spec, PointerType):
            return f"{self.c_type(type_spec)} {name}"
        return f"{self.c_type(type_spec)} {name}"

def compile_cct(source_code):
    tokens = lex(source_code)
    parser = Parser(tokens)
    ast = parser.parse()
    gen = CGenerator()
    gen.generate(ast)
    return '\n'.join(gen.code)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('input', help='Input .cct file')
    parser.add_argument('-o', '--output', default='output.c', help='Output C file')
    args = parser.parse_args()
    with open(args.input, 'r') as f: code = f.read()
    c_code = compile_cct(code)
    with open(args.output, 'w') as f: f.write(c_code)
    print(f'Compiled to {args.output}')

if __name__ == '__main__':
    main()