#!/usr/bin/env python3
"""Apply a known-good optimization to the target repo.

Handles multi-file optimizations by parsing the markdown code block format
(```python:filepath\ncode\n```) and applying changes to each file.

Uses codeflash's replace_functions_and_add_imports for each target file.

Configuration is read from /solution/config.json:
{
    "target_file": "inference/core/...",
    "function_name": "short_function_name"
}
"""

import ast
import json
import sys
from pathlib import Path

CONFIG = json.loads(Path("/solution/config.json").read_text(encoding="utf-8"))
TARGET_FILE = Path("/workspace/inference") / CONFIG["target_file"]
FUNCTION_NAME = CONFIG["function_name"]
RAW_OPTIMIZED = Path("/solution/optimized_code.py").read_text(encoding="utf-8")
PROJECT_ROOT = Path("/workspace/inference")


def parse_markdown_code_blocks(raw: str) -> list[tuple[str | None, str]]:
    """Parse ```python:filepath\\ncode\\n``` blocks. Returns [(filepath, code), ...]."""
    blocks = []
    current_lines = []
    current_path = None
    in_block = False

    for line in raw.splitlines():
        stripped = line.strip()
        if stripped.startswith("```") and not in_block:
            in_block = True
            current_lines = []
            if ":" in stripped:
                current_path = stripped.split(":", 1)[1].strip()
            else:
                current_path = None
        elif stripped == "```" and in_block:
            in_block = False
            blocks.append((current_path, "\n".join(current_lines)))
            current_lines = []
        elif in_block:
            current_lines.append(line)
        else:
            current_lines.append(line)

    # If no fences, treat whole thing as single block for TARGET_FILE
    if not blocks and current_lines:
        blocks.append((None, "\n".join(current_lines)))

    return blocks


def apply_to_file(file_path: Path, optimized_code: str, function_name: str) -> bool:
    """Apply optimized code to a single file using codeflash's production replacer.

    Uses replace_function_definitions_in_module which handles:
    - Function replacement
    - Module-level global assignments (constants, caches)
    - New helper functions
    - Import additions
    """
    if not file_path.exists():
        print(f"  File not found: {file_path}")
        return False

    try:
        ast.parse(optimized_code)
    except SyntaxError:
        print(f"  Optimized code has syntax error for {file_path}")
        return False

    # Collect qualified function names (ClassName.method for class methods, bare name for top-level)
    source = file_path.read_text(encoding="utf-8")

    def collect_qualified_names(code: str) -> set[str]:
        names = set()
        tree = ast.parse(code)
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                names.add(node.name)
            elif isinstance(node, ast.ClassDef):
                for child in ast.iter_child_nodes(node):
                    if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        names.add(f"{node.name}.{child.name}")
        return names

    opt_functions = collect_qualified_names(optimized_code)
    src_functions = collect_qualified_names(source)
    functions_to_replace = list(opt_functions & src_functions)
    if not functions_to_replace:
        print(f"  No matching functions found in {file_path}")
        return False

    try:
        from codeflash.languages.python.static_analysis.code_replacer import replace_function_definitions_in_module
        from codeflash.models.models import CodeStringsMarkdown

        # Wrap the optimized code in CodeStringsMarkdown format
        optimized_md = CodeStringsMarkdown.parse_markdown_code(
            f"```python:{file_path.relative_to(PROJECT_ROOT).as_posix()}\n{optimized_code}\n```"
        )

        did_update = replace_function_definitions_in_module(
            function_names=functions_to_replace,
            optimized_code=optimized_md,
            module_abspath=file_path,
            preexisting_objects=set(),
            project_root_path=PROJECT_ROOT,
        )
        if did_update:
            print(f"  Applied to {file_path.relative_to(PROJECT_ROOT)} (functions: {functions_to_replace})")
        else:
            print(f"  No changes for {file_path.relative_to(PROJECT_ROOT)}")
        return did_update
    except Exception as e:
        print(f"  Replacement failed for {file_path}: {e}")
        return False


def main() -> None:
    blocks = parse_markdown_code_blocks(RAW_OPTIMIZED)
    any_applied = False

    if len(blocks) == 1 and blocks[0][0] is None:
        # Single block, no file path — apply to target file
        _, code = blocks[0]
        if apply_to_file(TARGET_FILE, code, FUNCTION_NAME):
            any_applied = True
    else:
        # Multi-file: apply each block to its target file
        for file_path_str, code in blocks:
            if file_path_str:
                target = PROJECT_ROOT / file_path_str
            else:
                target = TARGET_FILE
            if apply_to_file(target, code, FUNCTION_NAME):
                any_applied = True

    if any_applied:
        # Verify syntax of all modified files
        for file_path_str, _ in blocks:
            target = PROJECT_ROOT / file_path_str if file_path_str else TARGET_FILE
            if target.exists():
                try:
                    ast.parse(target.read_text(encoding="utf-8"))
                except SyntaxError as e:
                    print(f"WARNING: Syntax error in {target}: {e}")
                    sys.exit(1)
        print("Syntax check passed")
    else:
        print("WARNING: No optimizations were applied")
        sys.exit(1)


if __name__ == "__main__":
    main()
