import subprocess
import re
import numpy as np
from pathlib import Path

class IRFeatureExtractor:
    FEATURE_NAMES = [
        "loop_depth_ratio",
        "memory_op_ratio",
        "branch_ratio",
        "arithmetic_ratio",
        "function_call_ratio",
        "phi_node_ratio",
        "bb_size_normalized",
        "num_functions_normalized",   # ★ 추가: 함수 개수
        "call_per_bb_normalized",     # ★ 추가: BB당 call 밀도
    ]

    def extract(self, bc_path: str) -> np.ndarray:
        ll_text = self._disassemble(bc_path)
        return self._parse_features(ll_text)

    def _disassemble(self, bc_path: str) -> str:
        result = subprocess.run(
            ["llvm-dis", bc_path, "-o", "-"],
            capture_output=True, text=True
        )
        if result.returncode != 0:
            raise RuntimeError(f"llvm-dis 실패: {result.stderr}")
        return result.stdout

    def _parse_features(self, ll_text: str) -> np.ndarray:
        lines = ll_text.split("\n")
        total_instr = 0
        mem_ops     = 0
        branches    = 0
        arith_ops   = 0
        call_ops    = 0
        phi_nodes   = 0
        bb_sizes    = []
        current_bb_size = 0
        loop_keywords   = 0
        num_functions   = 0

        for line in lines:
            stripped = line.strip()
            if not stripped or stripped.startswith(";"):
                continue
            # 함수 정의 카운트
            if re.match(r"^define ", stripped):
                num_functions += 1
                continue
            if re.match(r"^\w[\w.]*:$", stripped) or stripped.endswith(":"):
                if current_bb_size > 0:
                    bb_sizes.append(current_bb_size)
                current_bb_size = 0
                continue
            if re.search(r"\bcall\b", stripped):
                call_ops += 1
                total_instr += 1
                current_bb_size += 1
            elif any(op in stripped for op in ["load ", "store ", "getelementptr"]):
                mem_ops += 1
                total_instr += 1
                current_bb_size += 1
            elif stripped.startswith("br ") or stripped.startswith("switch "):
                branches += 1
                total_instr += 1
                current_bb_size += 1
            elif stripped.startswith("phi "):
                phi_nodes += 1
                total_instr += 1
                current_bb_size += 1
            elif any(stripped.startswith(op) for op in
                     ["%", "add ", "sub ", "mul ", "div ",
                      "fadd", "fsub", "fmul", "fdiv", "icmp", "fcmp"]):
                arith_ops += 1
                total_instr += 1
                current_bb_size += 1
            if "!llvm.loop" in stripped or "loop" in stripped.lower():
                loop_keywords += 1

        if current_bb_size > 0:
            bb_sizes.append(current_bb_size)

        safe_total  = max(total_instr, 1)
        num_bbs     = max(len(bb_sizes), 1)
        avg_bb_size = np.mean(bb_sizes) if bb_sizes else 1.0

        raw = np.array([
            min(loop_keywords / 10.0, 1.0),
            mem_ops      / safe_total,
            branches     / safe_total,
            arith_ops    / safe_total,
            call_ops     / safe_total,
            phi_nodes    / safe_total,
            min(avg_bb_size / 50.0, 1.0),
            min(num_functions / 20.0, 1.0),     # ★ 20개 함수 = 1.0
            min(call_ops / num_bbs / 2.0, 1.0), # ★ BB당 call 2개 = 1.0
        ], dtype=np.float32)
        return np.clip(raw, 0.0, 1.0)

if __name__ == "__main__":
    extractor = IRFeatureExtractor()
    features = extractor.extract("test_loop.bc")
    print("=== IR 특징 벡터 (9차원) ===")
    for name, val in zip(IRFeatureExtractor.FEATURE_NAMES, features):
        bar = "█" * int(val * 20)
        print(f"  {name:<30} {val:.4f}  {bar}")

def extract_features(bc_path: str) -> np.ndarray:
    return IRFeatureExtractor().extract(bc_path)