from __future__ import annotations

import ast
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Any
from typing import DefaultDict
from typing import Tuple
from typing import cast

from asttokens.util import Token
from executing.executing import EnhancedAST

from inline_snapshot._external._external_location import Location
from inline_snapshot._external._find_external import ensure_import
from inline_snapshot._source_file import SourceFile

from ._rewrite_code import ChangeRecorder
from ._rewrite_code import end_of
from ._rewrite_code import start_of

if TYPE_CHECKING:
    from inline_snapshot._external._format._protocol import Format


class ChangeBase:
    flag: str

    def rich_diff(self):
        raise NotImplementedError

    def apply_external_changes(self):
        raise NotImplementedError

    def apply(self, recorder: ChangeRecorder):
        raise NotImplementedError


@dataclass()
class ExternalChange(ChangeBase):
    flag: str

    new_file: Path

    old_location: Location
    new_location: Location

    format: Format

    def rich_diff(self):

        title = str(self.new_location)
        if title != (old_name := str(self.old_location)) and self.old_location.stem:
            title = f"{old_name} -> {title}"

        if (
            self.old_location.exists()
            and self.old_location.suffix == self.new_location.suffix
        ):
            with self.old_location.load() as old_file:
                return title, self.format.rich_diff(old_file, self.new_file)
        else:
            return title, self.format.rich_show(self.new_file)

    def apply_external_changes(self):
        self.new_location.store(self.new_file)

    def apply(self, recorder: ChangeRecorder):
        pass


@dataclass()
class ExternalRemove(ChangeBase):
    flag: str

    old_location: Location

    def rich_diff(self):

        title = f"delete {self.old_location}"

        from inline_snapshot._external._format._protocol import (
            get_format_handler_from_suffix,
        )

        with self.old_location.load() as old_file:
            assert self.old_location.suffix
            format = get_format_handler_from_suffix(self.old_location.suffix)
            return title, format.rich_show(old_file)

    def apply_external_changes(self):
        self.old_location.delete()

    def apply(self, recorder: ChangeRecorder):
        pass


@dataclass()
class Change(ChangeBase):
    flag: str
    file: SourceFile

    def rich_diff(self):
        return None

    @property
    def filename(self):
        return self.file.filename

    def apply(self, recorder: ChangeRecorder):
        raise NotImplementedError()

    def apply_external_changes(self):
        pass


@dataclass()
class RequiredImport(Change):
    module: str
    name: str | None = None


@dataclass()
class Delete(Change):
    node: ast.AST | None
    old_value: Any


@dataclass()
class ListInsert(Change):
    node: ast.List | ast.Tuple
    position: int

    new_code: list[str]
    new_values: list[Any]

    def __post_init__(self):
        self.new_code = [self.file.format_expression(v) for v in self.new_code]


@dataclass()
class DictInsert(Change):
    node: ast.Dict
    position: int

    new_code: list[tuple[str, str]]
    new_values: list[tuple[Any, Any]]

    def __post_init__(self):
        self.new_code = [
            (self.file.format_expression(k), self.file.format_expression(v))
            for k, v in self.new_code
        ]


@dataclass()
class Replace(Change):
    node: ast.AST

    new_code: str
    old_value: Any
    new_value: Any

    def apply(self, recorder: ChangeRecorder):
        change = recorder.new_change()
        range = self.file.asttokens().get_text_positions(self.node, False)
        change.replace(range, self.new_code, filename=self.filename)

    def __post_init__(self):
        self.new_code = self.file.format_expression(self.new_code)


@dataclass()
class CallArg(Change):
    node: ast.Call | None
    arg_pos: int | None
    arg_name: str | None

    new_code: str
    new_value: Any

    def __post_init__(self):
        self.new_code = self.file.format_expression(self.new_code)


TokenRange = Tuple[Token, Token]


def brace_tokens(source, node) -> TokenRange:
    first_token, *_, end_token = source.asttokens().get_tokens(node)
    return first_token, end_token


def generic_sequence_update(
    source: SourceFile,
    parent: ast.List | ast.Tuple | ast.Dict | ast.Call,
    brace_tokens: TokenRange,
    parent_elements: list[TokenRange | None],
    to_insert: dict[int, list[str]],
    recorder: ChangeRecorder,
):
    rec = recorder.new_change()

    new_code = []
    deleted = False
    last_token, end_token = brace_tokens
    is_start = True
    elements = 0

    for index, entry in enumerate(parent_elements):
        if index in to_insert:
            new_code += to_insert[index]
        if entry is None:
            deleted = True
        else:
            first_token, new_last_token = entry
            elements += len(new_code) + 1

            if deleted or new_code:

                code = ""
                if new_code:
                    code = ", ".join(new_code) + ", "
                if not is_start:
                    code = ", " + code

                rec.replace(
                    (end_of(last_token), start_of(first_token)),
                    code,
                    filename=source.filename,
                )
            new_code = []
            deleted = False
            last_token = new_last_token
            is_start = False

    for i, insert_elements in sorted(to_insert.items()):
        if i >= len(parent_elements):
            new_code += insert_elements
            elements += len(new_code)

    if new_code or deleted or elements == 1 or len(parent_elements) <= 1:
        code = ", ".join(new_code)
        if not is_start and code:
            code = ", " + code

        if elements == 1 and isinstance(parent, ast.Tuple):
            # trailing comma for tuples (1,)
            code += ","

        rec.replace(
            (end_of(last_token), start_of(end_token)),
            code,
            filename=source.filename,
        )


def apply_all(all_changes: list[ChangeBase], recorder: ChangeRecorder):
    by_parent: dict[EnhancedAST, list[Delete | DictInsert | ListInsert | CallArg]] = (
        defaultdict(list)
    )
    sources: dict[EnhancedAST, SourceFile] = {}

    # file -> module -> names
    imports_by_file: dict[str, dict[str, set]] = defaultdict(lambda: defaultdict(set))
    module_imports_by_file: dict[str, set] = defaultdict(set)

    for change in all_changes:
        if isinstance(change, Delete):
            node = cast(EnhancedAST, change.node).parent
            if isinstance(node, ast.keyword):
                node = node.parent
            by_parent[node].append(change)
            sources[node] = change.file

        elif isinstance(change, (DictInsert, ListInsert, CallArg)):
            node = cast(EnhancedAST, change.node)
            by_parent[node].append(change)
            sources[node] = change.file
        elif isinstance(change, RequiredImport):
            if change.name:
                imports_by_file[change.file.filename][change.module].add(change.name)
            else:
                module_imports_by_file[change.file.filename].add(change.module)
        else:
            change.apply(recorder)

    for filename in set(imports_by_file) | set(module_imports_by_file):
        imports = imports_by_file.get(filename, defaultdict(set))
        module_imports = module_imports_by_file.get(filename, set())
        ensure_import(filename, imports, module_imports, recorder)

    for parent, changes in by_parent.items():
        source = sources[parent]

        if isinstance(parent, (ast.List, ast.Tuple)):
            to_delete = {
                change.node for change in changes if isinstance(change, Delete)
            }
            to_insert = {
                change.position: change.new_code
                for change in changes
                if isinstance(change, ListInsert)
            }

            def list_token_range(entry):
                r = list(source.asttokens().get_tokens(entry))
                return r[0], r[-1]

            generic_sequence_update(
                source,
                parent,
                brace_tokens(source, parent),
                [None if e in to_delete else list_token_range(e) for e in parent.elts],
                to_insert,
                recorder,
            )

        elif isinstance(parent, ast.Call):
            to_delete = {
                change.node for change in changes if isinstance(change, Delete)
            }
            atok = source.asttokens()

            def arg_token_range(node):
                if isinstance(node.parent, ast.keyword):
                    node = node.parent
                r = list(atok.get_tokens(node))
                return r[0], r[-1]

            braces_left = atok.next_token(list(atok.get_tokens(parent.func))[-1])
            assert braces_left.string == "("
            braces_right = list(atok.get_tokens(parent))[-1]
            assert braces_right.string == ")"

            to_insert = DefaultDict(list)

            other_call_args: list[CallArg] = []
            call_args: list[CallArg] = []

            for change in changes:
                if isinstance(change, CallArg):
                    call_args.append(change)

            max_pos = len(parent.args)

            for change in sorted(call_args, key=lambda arg: arg.arg_pos or 0):
                if change.arg_name is None:
                    assert change.arg_pos is not None
                    to_insert[min(change.arg_pos, max_pos)].append(change.new_code)
                else:
                    other_call_args.append(change)

            end_pos = (
                max([len(parent.args) + len(parent.keywords), *to_insert.keys()]) + 1
            )

            for change in other_call_args:
                position = change.arg_pos if change.arg_pos is not None else end_pos
                to_insert[position].append(f"{change.arg_name}={change.new_code}")

            generic_sequence_update(
                source,
                parent,
                (braces_left, braces_right),
                [
                    None if e in to_delete else arg_token_range(e)
                    for e in parent.args + [kw.value for kw in parent.keywords]
                ],
                to_insert,
                recorder,
            )

        elif isinstance(parent, ast.Dict):
            to_delete = {
                change.node for change in changes if isinstance(change, Delete)
            }
            to_insert = {
                change.position: [f"{key}: {value}" for key, value in change.new_code]
                for change in changes
                if isinstance(change, DictInsert)
            }

            def dict_token_range(key, value):
                return (
                    list(source.asttokens().get_tokens(key))[0],
                    list(source.asttokens().get_tokens(value))[-1],
                )

            generic_sequence_update(
                source,
                parent,
                brace_tokens(source, parent),
                [
                    None if value in to_delete else dict_token_range(key, value)
                    for key, value in zip(parent.keys, parent.values)
                ],
                to_insert,
                recorder,
            )

        else:
            assert False, parent
