# Copyright 2025 the original author or authors.
# <p>
# Licensed under the Moderne Source Available License (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# <p>
# https://docs.moderne.io/licensing/moderne-source-available-license
# <p>
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""AddImport visitor for Python import handling."""

from dataclasses import dataclass
from enum import Enum
from typing import Optional

from rewrite import random_id
from rewrite.java import J
from rewrite.java.support_types import JContainer, JLeftPadded, JRightPadded
from rewrite.java.tree import Empty, FieldAccess, Identifier, Import, Space
from rewrite.markers import Markers
from rewrite.python.import_utils import get_qualid_name, get_name_string, get_alias_name, pad_right
from rewrite.python.tree import CompilationUnit, MultiImport
from rewrite.python.visitor import PythonVisitor


class ImportStyle(Enum):
    """Python import styles."""
    DIRECT = "direct"  # import module
    FROM = "from"  # from module import name


@dataclass
class AddImportOptions:
    """Options for adding an import.

    Attributes:
        module: The module to import from (e.g., 'os.path', 'typing')
        name: The specific name to import. If None, imports the module directly.
              Special values:
              - '*': Adds a star import (from module import *)
        alias: Optional alias for the imported name (e.g., 'np' for numpy)
        only_if_referenced: If True, only add the import if the name is used (default: True)
    """
    module: str
    name: Optional[str] = None
    alias: Optional[str] = None
    only_if_referenced: bool = True


def maybe_add_import(visitor: PythonVisitor, options: AddImportOptions) -> None:
    """Register an AddImport visitor to add an import statement.

    This function schedules the import to be added after the current visit completes.

    Args:
        visitor: The visitor to register the import addition with
        options: Configuration for the import to add

    Examples:
        # Add: import os
        maybe_add_import(visitor, AddImportOptions(module='os'))

        # Add: from os import path
        maybe_add_import(visitor, AddImportOptions(module='os', name='path'))

        # Add: from os.path import join
        maybe_add_import(visitor, AddImportOptions(module='os.path', name='join'))

        # Add: from os.path import join as pjoin
        maybe_add_import(visitor, AddImportOptions(module='os.path', name='join', alias='pjoin'))

        # Add: import numpy as np
        maybe_add_import(visitor, AddImportOptions(module='numpy', alias='np'))

        # Add: from typing import *
        maybe_add_import(visitor, AddImportOptions(module='typing', name='*'))
    """
    # Check for duplicate registrations
    if visitor._after_visit is None:
        visitor._after_visit = []

    for v in visitor._after_visit:
        if isinstance(v, AddImport):
            if (v.module == options.module and
                v.name == options.name and
                v.alias == options.alias):
                return  # Already registered

    visitor._after_visit.append(AddImport(options))


class AddImport(PythonVisitor):
    """Visitor that adds an import statement to a Python file.

    This visitor handles:
    - Direct imports: import module [as alias]
    - From imports: from module import name [as alias]
    - Star imports: from module import *
    - Merging into existing imports from the same module
    """

    def __init__(self, options: AddImportOptions):
        super().__init__()
        self.module = options.module
        self.name = options.name
        self.alias = options.alias
        self.only_if_referenced = options.only_if_referenced

    def visit_compilation_unit(self, cu: CompilationUnit, p) -> J:
        # Check if import already exists
        if self._import_exists(cu):
            return cu

        # If only_if_referenced is True, check if the identifier is used
        if self.only_if_referenced:
            if not self._is_referenced(cu):
                return cu

        # Try to merge into an existing import from the same module
        merged = self._try_merge_into_existing(cu)
        if merged is not cu:
            return merged

        # Add new import statement
        return self._add_import(cu)

    def _import_exists(self, cu: CompilationUnit) -> bool:
        """Check if the import already exists."""
        for stmt in cu.statements:
            if isinstance(stmt, Import) and not isinstance(stmt, MultiImport):
                if self.name is None and self._import_name_matches(stmt, self.module, self.alias):
                    return True
            elif isinstance(stmt, MultiImport):
                if self._multi_import_matches(stmt):
                    return True
        return False

    def _multi_import_matches(self, multi: MultiImport) -> bool:
        """Check if a MultiImport matches the import we want to add."""
        if self.name is None:
            # We want: import module [as alias]
            if multi.from_ is not None:
                return False  # This is a "from" import
            for imp in multi.names:
                if self._import_name_matches(imp, self.module, self.alias):
                    return True
        else:
            # We want: from module import name [as alias]
            if multi.from_ is None:
                return False  # This is not a "from" import
            from_name = get_name_string(multi.from_)
            if from_name != self.module:
                return False
            for imp in multi.names:
                if self._import_name_matches(imp, self.name, self.alias):
                    return True
        return False

    def _import_name_matches(self, imp: Import, name: str, alias: Optional[str]) -> bool:
        """Check if an Import matches the given name and alias."""
        import_name = get_qualid_name(imp.qualid)
        import_alias = get_alias_name(imp)

        if import_name != name:
            return False
        if alias is not None and import_alias != alias:
            return False
        if alias is None and import_alias is not None:
            return False
        return True

    def _is_referenced(self, cu: CompilationUnit) -> bool:
        """Check if the identifier we're importing is actually used."""
        target_name = self.alias or self.name or self.module.split('.')[-1]

        class ReferenceChecker(PythonVisitor):
            def __init__(self):
                super().__init__()
                self.found = False

            def visit_identifier(self, ident: Identifier, p) -> J:
                if ident.simple_name == target_name:
                    self.found = True
                return ident

        checker = ReferenceChecker()
        checker.visit(cu, None)
        return checker.found

    def _try_merge_into_existing(self, cu: CompilationUnit) -> CompilationUnit:
        """Try to merge the new import into an existing import from the same module."""
        if self.name is None:
            # Direct import - cannot merge
            return cu

        # Look for existing "from module import ..." statements
        for i, stmt in enumerate(cu.statements):
            if not isinstance(stmt, MultiImport):
                continue
            if stmt.from_ is None:
                continue
            from_name = get_name_string(stmt.from_)
            if from_name != self.module:
                continue

            # Found an existing import from the same module - add our name
            new_import = self._create_import_element(self.name, self.alias)
            # Give the new import a space prefix (space after comma)
            new_import = new_import.replace(prefix=Space.SINGLE_SPACE)

            # Preserve existing padded elements; append new one with proper padding
            existing_padded = list(stmt.padding.names.padding.elements)
            # Move last element's .after (typically '\n') to the new element
            if existing_padded:
                last = existing_padded[-1]
                existing_padded[-1] = last.replace(_after=Space.EMPTY)
                new_padded_import = JRightPadded(new_import, last.after, Markers.EMPTY)
            else:
                new_padded_import = JRightPadded(new_import, Space.EMPTY, Markers.EMPTY)
            existing_padded.append(new_padded_import)

            # Recreate the MultiImport with the new names
            new_multi = MultiImport(
                stmt.id,
                stmt.prefix,
                stmt.markers,
                stmt.padding.from_,
                stmt.parenthesized,
                JContainer(
                    stmt.padding.names.before,
                    existing_padded,
                    stmt.padding.names.markers
                )
            )

            # Replace the statement at the padding level
            padded_stmts = list(cu.padding.statements)
            old_padded = padded_stmts[i]
            padded_stmts[i] = JRightPadded(new_multi, old_padded.after, old_padded.markers)
            return cu.padding.replace(_statements=padded_stmts)

        return cu

    def _add_import(self, cu: CompilationUnit) -> CompilationUnit:
        """Add a new import statement to the compilation unit."""
        new_import = self._create_multi_import()

        # Find insertion point (after existing imports)
        insert_idx = 0
        padded_stmts = list(cu.padding.statements)
        for i, padded in enumerate(padded_stmts):
            if isinstance(padded.element, (Import, MultiImport)):
                insert_idx = i + 1
            elif insert_idx > 0:
                break  # Stop after we've passed the import section

        # Insert the new import at the padding level
        if insert_idx == 0 and padded_stmts:
            first = padded_stmts[0]
            first_prefix = first.element.prefix
            if first_prefix.whitespace and first_prefix.whitespace.startswith('\n'):
                # First statement has a leading newline (e.g., a prior import was
                # removed). The new import becomes the true first statement with
                # empty prefix; the existing statement keeps its '\n'.
                new_padded = JRightPadded(new_import.replace(prefix=Space.EMPTY), Space.EMPTY, Markers.EMPTY)
            else:
                # Normal case: new import takes the first statement's prefix,
                # and the first statement gets '\n' as separator.
                new_padded = JRightPadded(new_import.replace(prefix=first_prefix), Space.EMPTY, Markers.EMPTY)
                padded_stmts[0] = JRightPadded(
                    first.element.replace(prefix=Space([], '\n')),
                    first.after, first.markers
                )
            padded_stmts.insert(insert_idx, new_padded)
        else:
            # Inserting after existing imports. The new import needs '\n' as
            # its prefix since each statement's prefix carries the newline
            # that separates it from the preceding statement.
            new_padded = JRightPadded(new_import.replace(prefix=Space([], '\n')), Space.EMPTY, Markers.EMPTY)
            padded_stmts.insert(insert_idx, new_padded)

        return cu.padding.replace(_statements=padded_stmts)

    def _create_multi_import(self) -> MultiImport:
        """Create a new MultiImport statement."""
        if self.name is None:
            # Direct import: import module [as alias]
            import_elem = self._create_import_element(self.module, self.alias)
            return MultiImport(
                random_id(),
                Space([], '\n'),
                Markers.EMPTY,
                None,  # No 'from'
                False,  # Not parenthesized
                JContainer(
                    Space.SINGLE_SPACE,
                    [pad_right(import_elem)],
                    Markers.EMPTY
                )
            )
        else:
            # From import: from module import name [as alias]
            from_name = self._create_module_name(self.module)
            # Add space prefix (the space between 'from' and module name)
            from_name = from_name.replace(prefix=Space.SINGLE_SPACE)
            import_elem = self._create_import_element(self.name, self.alias)
            return MultiImport(
                random_id(),
                Space([], '\n'),
                Markers.EMPTY,
                JRightPadded(from_name, Space.SINGLE_SPACE, Markers.EMPTY),
                False,  # Not parenthesized
                JContainer(
                    Space.SINGLE_SPACE,
                    [pad_right(import_elem)],
                    Markers.EMPTY
                )
            )

    def _create_import_element(self, name: str, alias: Optional[str]) -> Import:
        """Create an Import element."""
        qualid = self._create_qualified_name(name)
        alias_left_padded = None
        if alias:
            alias_ident = Identifier(
                random_id(),
                Space.SINGLE_SPACE,
                Markers.EMPTY,
                [],
                alias,
                None,
                None
            )
            alias_left_padded = JLeftPadded(
                Space.SINGLE_SPACE,
                alias_ident,
                Markers.EMPTY
            )

        return Import(
            random_id(),
            Space.EMPTY,
            Markers.EMPTY,
            JLeftPadded(Space.EMPTY, False, Markers.EMPTY),
            qualid,
            alias_left_padded
        )

    def _create_qualified_name(self, name: str) -> FieldAccess:
        """Create a FieldAccess for a qualified name."""
        parts = name.split('.')
        if len(parts) == 1:
            return FieldAccess(
                random_id(),
                Space.EMPTY,
                Markers.EMPTY,
                Empty(random_id(), Space.EMPTY, Markers.EMPTY),
                JLeftPadded(
                    Space.EMPTY,
                    Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None),
                    Markers.EMPTY
                ),
                None
            )

        # Build nested FieldAccess for qualified names like "os.path"
        # Start with the first part as an Identifier target
        result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None)
        # Wrap remaining parts as FieldAccess nodes
        for part in parts[1:]:
            result = FieldAccess(
                random_id(),
                Space.EMPTY,
                Markers.EMPTY,
                result,
                JLeftPadded(
                    Space.EMPTY,
                    Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None),
                    Markers.EMPTY
                ),
                None
            )
        # For multi-part names, result is already a FieldAccess.
        # Wrap single-part Identifier in FieldAccess(Empty, name) for consistency
        # (but single-part is handled above, so this shouldn't happen)
        assert isinstance(result, FieldAccess)
        return result

    def _create_module_name(self, name: str) -> J:
        """Create a name tree for use as a 'from' module name.

        Single-part names become Identifier; multi-part become FieldAccess.
        Unlike _create_qualified_name, this does NOT wrap single parts in
        FieldAccess(Empty, name) since the 'from' printer has no special
        handling for Empty targets.
        """
        parts = name.split('.')
        result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None)
        for part in parts[1:]:
            result = FieldAccess(
                random_id(),
                Space.EMPTY,
                Markers.EMPTY,
                result,
                JLeftPadded(
                    Space.EMPTY,
                    Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None),
                    Markers.EMPTY
                ),
                None
            )
        return result

