import ast
import logging
import sys
from pathlib import Path
from typing import TYPE_CHECKING, List

import black


logger = logging.getLogger(__package__)

PATCHED_BLACK_NAME = f'blacked{black.__version__.replace(".", "_")}'


def log_changed(original_line, new_line):
    logger.debug(f'Changed:  {original_line.strip()}  to  {new_line.strip()}')


class CacheDirVisitor(ast.NodeVisitor):
    def __init__(self, source_lines: List[str]):
        self.source_lines = source_lines
        self.visited = False

    def visit_Call(self, node: ast.Call) -> None:
        if len(node.args) == 1 and isinstance(node.args[0], ast.Call):
            call = node.args[0]
            if getattr(call.func, 'id', ...) == 'user_cache_dir':
                line_index = node.lineno - 1
                original_line = self.source_lines[line_index]
                new_line = original_line.replace('"black"', '"blacked"')
                self.source_lines[line_index] = new_line
                self.visited = True
                log_changed(original_line, new_line)


class NormalizeStringQuotesVisitor(ast.NodeVisitor):
    def __init__(self, source_lines: List[str]):
        self.str_visitor = QuoteStringsVisitor(source_lines)
        self.visited = False

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        if node.name == 'normalize_string_quotes':
            self.str_visitor.visit(node)
            self.visited = True


class QuoteStringsVisitor(ast.NodeVisitor):
    def __init__(self, source_lines: List[str]):
        self.source_lines = source_lines

    def visit_Str(self, node: ast.Str) -> None:
        line_index = node.lineno - 1
        original_line = self.source_lines[line_index]
        string = node.s
        if string == '"':
            new_line = original_line.replace("'\"'", '"\'"')
        elif string == "'":
            new_line = original_line.replace('"\'"', "'\"'")
        else:
            return

        self.source_lines[line_index] = new_line
        log_changed(original_line, new_line)


def get_patched_black():
    """
    Will try to import patched black for currently installed black version
    If it fails, will create patched black for current black version
    """
    if TYPE_CHECKING:
        return black

    try:
        exec(f'from . import {PATCHED_BLACK_NAME}')
    except ImportError:
        logger.info(f'blacked {black.__version__} not found...')
        rewrite_black()
        exec(f'from . import {PATCHED_BLACK_NAME}')

    return locals()[PATCHED_BLACK_NAME]


def rewrite_black():
    logger.info(f'Creating patched black {black.__version__}...')
    black_code = Path(black.__file__).read_text()
    black_lines = black_code.split('\n')
    black_ast = ast.parse(black_code)

    for visitor_cls in NormalizeStringQuotesVisitor, CacheDirVisitor:
        visitor = visitor_cls(black_lines)
        visitor.visit(black_ast)
        if not visitor.visited:
            logger.warning(f'{visitor_cls.__name__} did not change any code...')

    new_black_module = Path(__file__).parent / (PATCHED_BLACK_NAME + '.py')

    # I used to change the AST itself and then used astunparse to generate new source code,
    # however, the code generated by astunparse produced strange behavior
    # (although ast.unparse from python 3.9 worked well, there's no equivalent backport AFAIK)
    new_black_module.write_text('\n'.join(black_lines))
    logger.info('Done!')


def filter_blacked_args():
    side_effects = {
        '--debug': lambda: logging.basicConfig(level=logging.DEBUG),
        '--force-patch': rewrite_black,
    }
    popped = 0
    for i, arg in enumerate(sys.argv.copy()):
        side_effect = side_effects.pop(arg, None)
        if side_effect is not None:
            sys.argv.pop(i - popped)
            popped += 1
            side_effect()


def main():
    filter_blacked_args()
    patched_black = get_patched_black()
    logging.disable(logging.DEBUG)

    if hasattr(patched_black, 'patched_main'):
        patched_black.patched_main()
    else:  # older versions
        patched_black.main()


if __name__ == '__main__':
    main()
