from __future__ import annotations

import argparse
import io
import re
import sys
import tokenize
from typing import Sequence


def handle_match(token_text: str, *, never: bool) -> str:
    if '"""' in token_text or "'''" in token_text:
        return token_text

    start_quote_re = re.compile('^[a-zA-Z]*"') if never else re.compile("^[a-zA-Z]*'")

    match = start_quote_re.match(token_text)
    if match is not None:
        meat = token_text[match.end() : -1]
        if '"' in meat or "'" in meat:
            return token_text
        else:
            if never:
                return match.group().replace('"', "'") + meat + "'"
            return match.group().replace("'", '"') + meat + '"'
    else:
        return token_text


def get_line_offsets_by_line_no(src: str) -> list[int]:
    # Padded so we can index with line number
    offsets = [-1, 0]
    for line in src.splitlines(True):
        offsets.append(offsets[-1] + len(line))
    return offsets


def fix_strings(filename: str, *, never: bool) -> int:
    with open(filename, encoding="UTF-8", newline="") as f:
        contents = f.read()
    line_offsets = get_line_offsets_by_line_no(contents)

    # Basically a mutable string
    splitcontents = list(contents)

    # Iterate in reverse so the offsets are always correct
    tokens_l = list(tokenize.generate_tokens(io.StringIO(contents).readline))
    tokens = reversed(tokens_l)

    inside_f_string = False

    for token_type, token_text, (srow, scol), (erow, ecol), _ in tokens:
        # Single quotes inside f-strings are only an issue in python 3.12
        if sys.version_info >= (3, 12):  # pragma: no cover
            if token_type == tokenize.FSTRING_END:
                inside_f_string = True
            if token_type == tokenize.FSTRING_START:
                inside_f_string = False
        else:  # pragma: no cover
            pass
        if token_type == tokenize.STRING:
            # do not touch any quotes inside of an f-string
            if inside_f_string:  # pragma: no cover can only be true in py312
                continue
            new_text = handle_match(token_text, never=never)
            splitcontents[line_offsets[srow] + scol : line_offsets[erow] + ecol] = (
                new_text
            )

    new_contents = "".join(splitcontents)
    if contents != new_contents:
        with open(filename, "w", encoding="UTF-8", newline="") as f:
            f.write(new_contents)
        return 1
    else:
        return 0


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("filenames", nargs="*", help="Filenames to fix")
    parser.add_argument("--never", action="store_true")
    args = parser.parse_args(argv)

    retv = 0

    for filename in args.filenames:
        return_value = fix_strings(filename, never=args.never)
        if return_value != 0:
            print(f"Fixing strings in {filename}")
        retv |= return_value

    return retv


if __name__ == "__main__":
    raise SystemExit(main())
