"""
transpiler/profiling_optimizer.py — Automatic optimization based on profiling data.
"""

from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Dict, Optional
from transpiler.middleend import ir
from transpiler.infrastructure.flags import CompilerFlags

class ProfilingOptimizer:
    def __init__(self, flags: CompilerFlags):
        self.flags = flags
        self.profile_data: Dict[str, Any] = {}
        self._load_profile()

    def _load_profile(self):
        if not self.flags.profile_path:
            return
            
        path = Path(self.flags.profile_path)
        if path.exists():
            try:
                self.profile_data = json.loads(path.read_text(encoding="utf-8"))
            except Exception as e:
                if self.flags.debug:
                    print(f"Warning: Failed to load profiling data from {path}: {e}")

    def optimize(self, module: ir.IRModule):
        if not self.profile_data:
            return

        funcs = self.profile_data.get("functions", {})
        if not funcs:
            return

        threshold = self.flags.auto_native_threshold
        
        # Walk IR and identify functions
        for stmt in module.body:
            if isinstance(stmt, ir.IRFunctionDef):
                self._optimize_function(stmt, funcs, threshold)
            elif isinstance(stmt, ir.IRClassDef):
                for member in stmt.body:
                    if isinstance(member, ir.IRFunctionDef):
                        # For classes, we currently use simple naming (MethodName)
                        # In the future, we might want Class.Method
                        self._optimize_function(member, funcs, threshold)

    def _optimize_function(self, node: ir.IRFunctionDef, funcs: Dict[str, Any], threshold: float):
        # Check by name
        stats = funcs.get(node.name)
        if stats:
            avg_ms = stats.get("avg_ms", 0.0)
            if avg_ms >= threshold:
                if "native" not in node.decorators:
                    node.decorators.append("native")
                    if self.flags.explain:
                        print(f"  [Opt] Injected @native for '{node.name}' (avg {avg_ms:.4f}ms >= threshold {threshold}ms)")
