import ast
import sys
from pathlib import Path

ROOT_DIR = Path(__file__).parent.parent

class IndentationFixer(ast.NodeVisitor):
    def __init__(self, source_lines: list[str]):
        self.source_lines = source_lines
        self.result_lines: list[str] = []
        self._pending_lines: list[tuple[int, str]] = []
        self._current_depth = 0

    def fix(self) -> list[str]:
        self._process_pending(0)
        self.visit_Module(ast.parse("".join(self.source_lines)))
        return self.result_lines

    def _get_indent(self, lineno: int) -> int:
        if lineno <= 0 or lineno > len(self.source_lines):
            return 0
        line = self.source_lines[lineno - 1]
        stripped = line.lstrip()
        return len(line) - len(stripped)

    def _is_docstring_or_comment(self, line: str) -> bool:
        stripped = line.lstrip()
        if stripped.startswith('#'):
            return True
        if stripped.startswith('"""') or stripped.startswith("'''"):
            return True
        return False

    def _process_pending(self, target_depth: int):
        while self._pending_lines:
            line_no, line = self._pending_lines[0]
            stripped = line.lstrip()
            actual = len(line) - len(stripped)
            expected = target_depth

            if actual == expected:
                self.result_lines.append(line)
                self._pending_lines.pop(0)
            elif actual < expected:
                break
            elif actual > expected:
                self.result_lines.append(" " * expected + stripped)
                self._pending_lines.pop(0)

    def visit_Module(self, node: ast.Module):
        for child in node.body:
            self._walk_node(child)
        self.generic_visit(node)

    def _walk_node(self, node: ast.AST):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            lineno = node.lineno
            for i, (old_no, old_line) in enumerate(self._pending_lines):
                if old_no == lineno:
                    self._pending_lines.pop(i)
                    break

            actual = self._get_indent(lineno)
            expected = self._current_depth
            line = self.source_lines[lineno - 1]
            stripped = line.lstrip()

            if actual != expected:
                self.result_lines.append(" " * expected + stripped)
            else:
                self.result_lines.append(line)

            self._current_depth += 1
            for child in node.body:
                self._walk_node(child)
            self._current_depth -= 1

        elif isinstance(node, ast.If):
            lineno = node.lineno
            line = self.source_lines[lineno - 1]
            stripped = line.lstrip()
            actual = self._get_indent(lineno)
            expected = self._current_depth

            if actual != expected:
                self.result_lines.append(" " * expected + stripped)
            else:
                self.result_lines.append(line)

            self._current_depth += 1
            for child in node.body:
                self._walk_node(child)
            self._current_depth -= 1

            if node.orelse:
                self._current_depth += 1
                self._walk_node(node.orelse)
                self._current_depth -= 1

        elif isinstance(node, (ast.For, ast.While, ast.With)):
            lineno = node.lineno
            line = self.source_lines[lineno - 1]
            stripped = line.lstrip()
            actual = self._get_indent(lineno)
            expected = self._current_depth

            if actual != expected:
                self.result_lines.append(" " * expected + stripped)
            else:
                self.result_lines.append(line)

            self._current_depth += 1
            for child in node.body:
                self._walk_node(child)
            self._current_depth -= 1

            if isinstance(node, ast.For) and node.orelse:
                self._current_depth += 1
                for child in node.orelse:
                    self._walk_node(child)
                self._current_depth -= 1

        elif isinstance(node, ast.Try):
            for child in node.body:
                self._walk_node(child)
            for handler in node.handlers:
                self._current_depth += 1
                for child in handler.body:
                    self._walk_node(child)
                self._current_depth -= 1
            if node.orelse:
                self._current_depth += 1
                for child in node.orelse:
                    self._walk_node(child)
                self._current_depth -= 1
            if node.finalbody:
                self._current_depth += 1
                for child in node.finalbody:
                    self._walk_node(child)
                self._current_depth -= 1

        else:
            self.generic_visit(node)

def fix_file_ast(filepath: Path) -> tuple[bool, str]:
    try:
        with open(filepath, "r", encoding="utf-8", newline="") as f:
            source = f.read()

        source_lines = source.splitlines()
        tree = ast.parse(source, filename=str(filepath))

        fixer = IndentationFixer(source_lines)
        new_lines = fixer.fix()

        new_source = "\n".join(new_lines)
        if new_source == source:
            return False, "No changes"

        with open(filepath, "w", encoding="utf-8", newline="") as f:
            f.write(new_source)

        ast.parse(new_source)
        return True, "Fixed"

    except SyntaxError as e:
        return False, f"Syntax error: {e}"
    except Exception as e:
        return False, str(e)

def fix_file_simple(filepath: Path, base_indent: int = 4) -> tuple[bool, str]:
    try:
        with open(filepath, "r", encoding="utf-8", newline="") as f:
            lines = f.readlines()

        in_docstring = False
        new_lines = []
        changed = False

        for line in lines:
            stripped = line.lstrip()
            if not stripped:
                new_lines.append(line)
                continue

            if not in_docstring:
                if stripped.startswith('#'):
                    new_lines.append(line)
                    continue
                if '"""' in stripped or "'''" in stripped:
                    triple_pos = max(stripped.find('"""') if '"""' in stripped else 999,
                                  stripped.find("'''") if "'''" in stripped else 999)
                    if triple_pos == 0:
                        in_docstring = True
                    new_lines.append(line)
                    continue

            if in_docstring:
                if '"""' in stripped or "'''" in stripped:
                    in_docstring = False
                new_lines.append(line)
                continue

            leading = len(line) - len(stripped)
            if leading > 0:
                level = leading // base_indent
                new_leading = level
                if leading != new_leading:
                    new_lines.append(" " * new_leading + stripped + ("\n" if not line.endswith("\n") else ""))
                    changed = True
                else:
                    new_lines.append(line)
            else:
                new_lines.append(line)

        if changed:
            with open(filepath, "w", encoding="utf-8", newline="") as f:
                f.writelines(new_lines)
            return True, "Fixed"
        return False, "No changes needed"

    except Exception as e:
        return False, str(e)

def main():
    if len(sys.argv) > 1:
        filepath = Path(sys.argv[1])
        changed, msg = fix_file_simple(filepath)
        print(f"{filepath}: {msg}")
        return

    files_to_fix = [
        ("src/fuzzy_anchor.py", 4),
        ("src/patch_modal.py", 2),
        ("scripts/extract_symbols.py", 4),
        ("scripts/tasks/download_fonts.py", 4),
        ("tests/test_arch_boundary_phase1.py", 2),
        ("tests/test_arch_boundary_phase2.py", 2),
        ("tests/test_arch_boundary_phase3.py", 2),
        ("tests/test_external_editor.py", 4),
        ("tests/test_headless_service.py", 4),
        ("tests/test_history_manager.py", 4),
        ("tests/test_fuzzy_anchor.py", 4),
        ("tests/test_gemini_cli_adapter.py", 4),
        ("tests/test_ai_client_cli.py", 4),
        ("tests/test_api_events.py", 4),
        ("tests/test_context_composition_decoupled.py", 4),
        ("tests/test_context_composition_phase3.py", 4),
        ("tests/test_context_composition_phase4.py", 4),
        ("tests/test_diff_viewer.py", 2),
        ("tests/test_discussion_takes_gui.py", 4),
        ("tests/test_external_mcp_hitl.py", 4),
        ("tests/test_gui_discussion_tabs.py", 4),
        ("tests/test_gui_stress_performance.py", 4),
        ("tests/test_gui_updates.py", 2),
        ("tests/test_hot_reloader.py", 4),
        ("tests/test_mma_dashboard_refresh.py", 4),
        ("tests/test_mma_node_editor.py", 4),
        ("tests/test_mma_orchestration_gui.py", 4),
        ("tests/test_py_struct_tools.py", 4),
        ("tests/test_thinking_persistence.py", 4),
        ("tests/test_tier4_interceptor.py", 2),
        ("tests/test_tiered_aggregation.py", 4),
        ("tests/test_visual_orchestration.py", 4),
    ]

    for rel_path, base_indent in files_to_fix:
        filepath = ROOT_DIR / rel_path
        if filepath.exists():
            changed, msg = fix_file_simple(filepath, base_indent)
            print(f"{rel_path}: {msg}")

if __name__ == "__main__":
    main()