#!/usr/bin/env python3
"""
Apply Type Hints from suggestions JSON to source files.

Reads /tmp/type-hints-suggestions.json (generated by generate-type-hints.py),
applies typed signatures to source files using AST-based function lookup.

Safety: validates each file with py_compile after modification; reverts on failure.
"""

import ast
import json
import py_compile
import sys
from pathlib import Path

PROTECTED_PATTERNS = [
    ".github/workflows/",
    ".env",
    "ods/installers/",
    "ods/ods-cli",
    "ods/config/",
]


def is_protected(file_path: str) -> bool:
    """Check if a file path matches any protected patterns."""
    for pattern in PROTECTED_PATTERNS:
        if file_path.startswith(pattern) or f"/{pattern}" in file_path:
            return True
    return False


def find_function_line(source: str, function_name: str) -> int | None:
    """Use AST to find the actual line number of a function by name."""
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return None

    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if node.name == function_name:
                return node.lineno
    return None


def find_def_end_line(lines: list[str], start_idx: int) -> int:
    """Find the line index where a def statement ends (the line with the colon)."""
    depth = 0
    for i in range(start_idx, len(lines)):
        line = lines[i]
        depth += line.count("(") - line.count(")")
        if depth <= 0 and ":" in line:
            stripped = line.rstrip()
            if stripped.endswith(":"):
                return i
            colon_pos = stripped.rfind(":")
            if colon_pos >= 0:
                return i
    return start_idx


def get_existing_imports(source: str) -> set[str]:
    """Extract all existing import statements from source."""
    imports = set()
    for line in source.splitlines():
        stripped = line.strip()
        if stripped.startswith("import ") or stripped.startswith("from "):
            imports.add(stripped)
    return imports


def find_last_import_line(lines: list[str]) -> int:
    """Find the 0-based index of the last import line in the file."""
    last_import = -1
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("import ") or stripped.startswith("from "):
            last_import = i
    return last_import


def normalize_import(imp: str) -> list[str]:
    """Normalize an import statement."""
    return [imp.strip()]


def apply_type_hints(suggestions_path: str) -> dict:
    """Apply generated type hint suggestions to Python function signatures."""
    with open(suggestions_path, "r") as f:
        data = json.load(f)

    functions = data.get("functions_annotated", [])
    if not functions:
        print("No type hint suggestions to apply.")
        return {"files_modified": 0, "functions_applied": 0, "files_reverted": 0}

    by_file: dict[str, list] = {}
    for func in functions:
        file_path = func.get("file", "")
        if not file_path:
            continue
        if is_protected(file_path):
            print(f"  Skipping protected file: {file_path}")
            continue
        if not Path(file_path).exists():
            print(f"  Skipping missing file: {file_path}")
            continue
        by_file.setdefault(file_path, []).append(func)

    files_modified = 0
    functions_applied = 0
    files_reverted = 0

    for file_path, file_funcs in by_file.items():
        print(f"\nProcessing {file_path} ({len(file_funcs)} functions)...")

        original_content = Path(file_path).read_text()
        source = original_content
        lines = source.splitlines(keepends=True)

        located = []
        for func in file_funcs:
            actual_line = find_function_line(source, func["function"])
            if actual_line is None:
                print(
                    f"  Could not find function '{func['function']}' in AST, skipping"
                )
                continue
            located.append((actual_line, func))

        located.sort(key=lambda x: x[0], reverse=True)

        applied_in_file = 0
        for actual_line, func in located:
            typed_sig = func.get("typed_signature", "").strip()
            if not typed_sig:
                continue

            start_idx = actual_line - 1
            if start_idx >= len(lines) or start_idx < 0:
                continue

            end_idx = find_def_end_line(lines, start_idx)

            current_line = lines[start_idx]
            indent = current_line[: len(current_line) - len(current_line.lstrip())]

            current_stripped = current_line.lstrip()
            is_async = current_stripped.startswith("async ")

            typed_stripped = typed_sig.lstrip()
            if is_async and not typed_stripped.startswith("async "):
                typed_sig = "async " + typed_sig
            elif not is_async and typed_stripped.startswith("async "):
                typed_sig = typed_sig.replace("async ", "", 1)

            if not typed_sig.rstrip().endswith(":"):
                typed_sig = typed_sig.rstrip() + ":"

            new_line = indent + typed_sig.strip() + "\n"

            lines[start_idx : end_idx + 1] = [new_line]
            applied_in_file += 1
            print(f"  Applied: {func['function']} (line {actual_line})")

        if applied_in_file == 0:
            continue

        new_source = "".join(lines)
        existing_imports = get_existing_imports(new_source)

        imports_to_add = []
        for func in file_funcs:
            for imp in func.get("imports_needed", []):
                for normalized in normalize_import(imp):
                    if normalized not in existing_imports:
                        imports_to_add.append(normalized)
                        existing_imports.add(normalized)

        if imports_to_add:
            lines = new_source.splitlines(keepends=True)
            insert_idx = find_last_import_line(lines)
            if insert_idx >= 0:
                insert_point = insert_idx + 1
            else:
                insert_point = 0
                for i, line in enumerate(lines):
                    stripped = line.strip()
                    if (
                        stripped.startswith("#")
                        or stripped.startswith('"""')
                        or stripped.startswith("'''")
                        or not stripped
                    ):
                        insert_point = i + 1
                    else:
                        break

            import_lines = [imp + "\n" for imp in imports_to_add]
            lines[insert_point:insert_point] = import_lines
            new_source = "".join(lines)

        Path(file_path).write_text(new_source)

        try:
            py_compile.compile(file_path, doraise=True)
            files_modified += 1
            functions_applied += applied_in_file
            print(f"  Validated: {file_path} ({applied_in_file} functions)")
        except py_compile.PyCompileError as e:
            print(f"  REVERT: {file_path} failed compilation: {e}")
            Path(file_path).write_text(original_content)
            files_reverted += 1

    return {
        "files_modified": files_modified,
        "functions_applied": functions_applied,
        "files_reverted": files_reverted,
    }


def main() -> None:
    """Load and apply generated type hints."""
    suggestions_path = (
        sys.argv[1] if len(sys.argv) > 1 else "/tmp/type-hints-suggestions.json"
    )

    if not Path(suggestions_path).exists():
        print(f"Suggestions file not found: {suggestions_path}")
        print("No type hints to apply.")
        return

    print(f"Loading suggestions from {suggestions_path}...")
    result = apply_type_hints(suggestions_path)

    print("\n## Type Hints Application Summary\n")
    print(f"- **Files modified**: {result['files_modified']}")
    print(f"- **Functions updated**: {result['functions_applied']}")
    print(f"- **Files reverted** (compilation failed): {result['files_reverted']}")


if __name__ == "__main__":
    main()