"""Tests for add_missing_imports postprocessing."""

import libcst as cst

from core.languages.python.cst_utils import DefinitionRemover, file_path_to_module_path
from core.languages.python.testgen.postprocessing.add_missing_imports import (
    add_future_annotations_import,
    add_missing_imports,
)

# =============================================================================
# Tests for add_missing_imports
# =============================================================================


def test_add_missing_imports_basic() -> None:
    """Test basic import addition from source CST module."""
    source_code = """
class PreChunk:
    pass

class ChunkingOptions:
    pass

class Element:
    pass
"""

    test_code = """
from mymodule import _PreChunkAccumulator

def test_something():
    chunk = PreChunk()
    opts = ChunkingOptions()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "mymodule")

    assert "PreChunk" in result.code
    assert "ChunkingOptions" in result.code


def test_add_missing_imports_preserves_existing() -> None:
    """Test that existing imports are preserved."""
    source_code = """
class ClassA:
    pass

class ClassB:
    pass
"""

    test_code = """
from mymodule import ClassA

def test_something():
    a = ClassA()
    b = ClassB()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "mymodule")

    assert "ClassA" in result.code
    assert "ClassB" in result.code


def test_add_missing_imports_no_changes_needed() -> None:
    """Test when all imports are already present."""
    source_code = """
class ClassA:
    pass
"""

    test_code = """
from mymodule import ClassA

def test_something():
    a = ClassA()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "mymodule")

    assert "ClassA" in result.code


def test_add_missing_imports_handles_redefined_classes() -> None:
    """Test removal of locally redefined classes and import addition."""
    source_code = """
class PreChunk:
    def __init__(self, elements):
        self.elements = elements

class ChunkingOptions:
    def __init__(self):
        pass

class Element:
    def __init__(self, text):
        self.text = text
"""

    test_code = """
from unstructured.chunking.base import _PreChunkAccumulator

class Element:
    def __init__(self, text):
        self.text = text

class ChunkingOptions:
    pass

def test_will_fit():
    opts = ChunkingOptions()
    acc = _PreChunkAccumulator(opts)
    chunk = PreChunk([Element("test")])
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "unstructured.chunking.base")

    # Local class definitions should be removed
    assert "class Element" not in result.code
    assert "class ChunkingOptions" not in result.code

    # All three should now be imported
    assert "from unstructured.chunking.base import" in result.code
    assert "PreChunk" in result.code
    assert "Element" in result.code
    assert "ChunkingOptions" in result.code


def test_add_missing_imports_handles_private_symbols() -> None:
    """Test that private symbols get imported when referenced."""
    source_code = """
class Pre:
    pass

class _PreElementAccumulator:
    def __init__(self, parent):
        self.parent = parent

    @classmethod
    def reset_counter(cls):
        pass
"""

    test_code = """
from mymodule import Pre

def test_element_accum():
    p = Pre()
    _PreElementAccumulator.reset_counter()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "mymodule")

    assert "_PreElementAccumulator" in result.code
    assert "from mymodule import" in result.code


def test_add_missing_imports_handles_source_module_imports() -> None:
    """Test that symbols imported by source are imported from their original source."""
    source_code = """
from typing import List
from unstructured.documents.elements import ElementMetadata

exclude_metadata_keys = ("coordinates", "data_source")

def create_unstructured_weaviate_class(class_name: str = "UnstructuredDocument"):
    for name, annotation in ElementMetadata.__annotations__.items():
        pass
"""

    test_code = """
def test_large_scale():
    annotations = {}
    ElementMetadata.__annotations__ = annotations
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "unstructured.staging.weaviate")

    assert "ElementMetadata" in result.code
    assert "from unstructured.documents.elements import ElementMetadata" in result.code


def test_add_missing_imports_handles_both_defined_and_imported_symbols() -> None:
    """Test handling of both local definitions and imported symbols."""
    source_code = """
from external.module import ExternalClass

class LocalClass:
    pass

def local_func():
    pass
"""

    test_code = """
def test_something():
    local = LocalClass()
    external = ExternalClass()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "mymodule")

    assert "LocalClass" in result.code
    assert "ExternalClass" in result.code
    assert "from external.module import ExternalClass" in result.code
    assert "from mymodule import" in result.code


def test_add_missing_imports_handles_referenced_but_not_defined_symbols() -> None:
    """Test fallback for symbols referenced but not defined in incomplete source snippets."""
    source_code = """
_chunker_registry: dict[str, _ChunkerSpec] = {
    "basic": _ChunkerSpec(chunk_elements),
    "by_title": _ChunkerSpec(chunk_by_title),
}

def chunk(elements, strategy):
    spec = _chunker_registry.get(strategy)
    return spec.func(elements)
"""

    test_code = """
def test_chunker():
    registry = {"basic": _ChunkerSpec(lambda x: x)}
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "unstructured.chunking.dispatch")

    assert "_ChunkerSpec" in result.code
    assert "from unstructured.chunking.dispatch import" in result.code


def test_add_missing_imports_handles_local_redefinition_of_imported_symbol() -> None:
    """Test that redefined symbols are imported from their original source."""
    source_code = """
# file: typeagent/knowpro/fuzzyindex.py
from typeagent.aitools.vectorbase import VectorBase

class EmbeddingIndex:
    def __init__(self, settings):
        self._vector_base = VectorBase(settings)

    async def size(self) -> int:
        return len(self._vector_base)

# file: typeagent/knowpro/textlocindex.py
from typeagent.knowpro.fuzzyindex import EmbeddingIndex

class TextToTextLocationIndex:
    def __init__(self, settings):
        self._embedding_index = EmbeddingIndex(settings=settings)

    async def size(self) -> int:
        return await self._embedding_index.size()
"""

    test_code = """
import pytest
from typeagent.knowpro.textlocindex import TextToTextLocationIndex

class EmbeddingIndex:
    '''Locally redefined by LLM'''
    def __init__(self, settings, embeddings=None):
        pass

    async def size(self) -> int:
        return 0

@pytest.mark.asyncio
async def test_size():
    emb_index = EmbeddingIndex(settings=None, embeddings=[[1, 2, 3]])
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, cst.parse_module(source_code), "typeagent.knowpro.textlocindex")

    assert "class EmbeddingIndex" not in result.code
    assert "Locally redefined by LLM" not in result.code
    assert "from typeagent.knowpro.fuzzyindex import EmbeddingIndex" in result.code
    assert "from typeagent.knowpro.textlocindex import TextToTextLocationIndex" in result.code


# =============================================================================
# Tests for add_missing_imports (multi-context dict[str, cst.Module])
# =============================================================================


def test_add_missing_imports_multi_context_correct_module() -> None:
    """Test that symbols are imported from their defining module in multi-context."""
    source_code_blocks = {
        "skyvern/forge/sdk/core/skyvern_context.py": cst.parse_module("""
class SkyvernContext:
    def __init__(self):
        self.data = {}
"""),
        "skyvern/core/script_generations/skyvern_page.py": cst.parse_module("""
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext

class SkyvernPage:
    def __init__(self):
        self.context = SkyvernContext()

    def get_context(self):
        return self.context
"""),
    }

    test_code = """
import pytest
from skyvern.core.script_generations.skyvern_page import SkyvernPage

def test_get_context():
    page = SkyvernPage()
    ctx = SkyvernContext()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, source_code_blocks, "skyvern.core.script_generations.skyvern_page")

    assert "from skyvern.forge.sdk.core.skyvern_context import SkyvernContext" in result.code
    assert "from skyvern.core.script_generations.skyvern_page import SkyvernPage" in result.code


def test_add_missing_imports_multi_context_local_redefinition() -> None:
    """Test that redefined classes are removed and imported correctly in multi-context."""
    source_code_blocks = {
        "module_a/core.py": cst.parse_module("""
class CoreClass:
    def __init__(self, value):
        self.value = value
"""),
        "module_b/wrapper.py": cst.parse_module("""
from module_a.core import CoreClass

class Wrapper:
    def __init__(self):
        self.core = CoreClass(42)
"""),
    }

    test_code = """
import pytest
from module_b.wrapper import Wrapper

class CoreClass:
    '''Redefined by LLM'''
    def __init__(self, value):
        self.value = value

def test_wrapper():
    w = Wrapper()
    c = CoreClass(10)
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, source_code_blocks, "module_b.wrapper")

    assert "class CoreClass" not in result.code
    assert "Redefined by LLM" not in result.code
    assert "from module_a.core import CoreClass" in result.code


def test_add_missing_imports_multi_context_fallback_symbols() -> None:
    """Test fallback for referenced but not defined symbols in multi-context."""
    source_code_blocks = {
        "mymodule/handlers.py": cst.parse_module("""
_handlers: dict[str, _SpecialHandler] = {}

def get_handler(name):
    return _handlers.get(name)
""")
    }

    test_code = """
def test_handlers():
    handler = _SpecialHandler()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, source_code_blocks, "mymodule.handlers")

    assert "_SpecialHandler" in result.code
    assert "from mymodule.handlers import" in result.code


def test_add_missing_imports_multi_context_empty_blocks() -> None:
    """Test handling of empty source code blocks."""
    source_code_blocks: dict[str, cst.Module] = {}

    test_code = """
def test_something():
    x = SomeClass()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, source_code_blocks, "default.module")

    assert result.code == test_code


def test_add_missing_imports_multi_context_multiple_files_same_symbol() -> None:
    """Test that first definition wins when symbol is defined in multiple files."""
    source_code_blocks = {
        "package/v1/model.py": cst.parse_module("""
class Model:
    version = 1
"""),
        "package/v2/model.py": cst.parse_module("""
class Model:
    version = 2
"""),
    }

    test_code = """
def test_model():
    m = Model()
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, source_code_blocks, "default.module")

    assert "from package.v1.model import Model" in result.code


def test_add_missing_imports_multi_context_aliased_imports() -> None:
    """Test that aliased imports are handled correctly."""
    source_code_blocks = {
        "skyvern/forge/sdk/core/skyvern_context.py": cst.parse_module("""
class SkyvernContext:
    current = None

    def __init__(self):
        self.data = {}
"""),
        "skyvern/core/script_generations/skyvern_page.py": cst.parse_module("""
from skyvern.forge.sdk.core import skyvern_context as skyvern_context_mod

class SkyvernPage:
    def type(self, text: str):
        ctx = skyvern_context_mod.current()
        return ctx
"""),
    }

    test_code = """
import pytest
from skyvern.core.script_generations.skyvern_page import SkyvernPage

def test_type():
    skyvern_context_mod.current = lambda: "mock"
    page = SkyvernPage()
    result = page.type("hello")
"""
    module = cst.parse_module(test_code)
    result = add_missing_imports(module, source_code_blocks, "skyvern.core.script_generations.skyvern_page")

    assert "from skyvern.forge.sdk.core import skyvern_context as skyvern_context_mod" in result.code


def test_definition_remover_preserves_nested_functions() -> None:
    """Test that DefinitionRemover only removes top-level definitions, not nested ones.

    This is critical for test code that defines helper functions inside test functions,
    such as thread workers or callbacks.
    """
    code = """
def worker():
    # Top-level worker function - should be removed if in names_to_remove
    pass

def test_concurrent_calls():
    results = []

    def worker(idx):
        # Nested worker function - should NOT be removed even if 'worker' is in names_to_remove
        results.append(idx)

    threads = [worker(i) for i in range(10)]
"""
    module = cst.parse_module(code)
    # Even though 'worker' is in names_to_remove, only the top-level one should be removed
    remover = DefinitionRemover({"worker"})
    result = module.visit(remover)

    # Top-level worker should be removed
    assert result.code.count("def worker") == 1  # Only the nested one remains

    # The nested worker inside test_concurrent_calls should still be there
    assert "def worker(idx):" in result.code
    assert "results.append(idx)" in result.code

    # Only the top-level worker should be in removed_names
    assert remover.removed_names == {"worker"}


def test_definition_remover_preserves_nested_classes() -> None:
    """Test that DefinitionRemover preserves nested class definitions."""
    code = """
class Helper:
    # Top-level Helper class - should be removed if in names_to_remove
    pass

def test_something():
    class Helper:
        # Nested Helper class - should NOT be removed
        value = 42

    h = Helper()
"""
    module = cst.parse_module(code)
    remover = DefinitionRemover({"Helper"})
    result = module.visit(remover)

    # Top-level Helper should be removed
    # The nested Helper inside test_something should still be there
    assert result.code.count("class Helper") == 1
    assert "value = 42" in result.code
    assert remover.removed_names == {"Helper"}


# =============================================================================
# Tests for file_path_to_module_path
# =============================================================================


def test_file_path_to_module_path() -> None:
    """Test file path to module path conversion."""
    assert file_path_to_module_path("a/b/c.py") == "a.b.c"
    assert file_path_to_module_path("a/b/c") == "a.b.c"
    assert file_path_to_module_path("single.py") == "single"
    assert file_path_to_module_path("a\\b\\c.py") == "a.b.c"
    assert (
        file_path_to_module_path("skyvern/forge/sdk/core/skyvern_context.py")
        == "skyvern.forge.sdk.core.skyvern_context"
    )


# =============================================================================
# Tests for add_future_annotations_import
# =============================================================================


def test_add_future_annotations_import_adds_when_needed() -> None:
    """Test that import is added when undefined annotation names exist."""
    code = """from contextvars import ContextVar

_context: ContextVar[SkyvernContext | None] = ContextVar("context")
"""
    module = cst.parse_module(code)
    result = add_future_annotations_import(module)

    expected = """from __future__ import annotations
from contextvars import ContextVar

_context: ContextVar[SkyvernContext | None] = ContextVar("context")
"""
    assert result.code == expected


def test_add_future_annotations_import_skips_when_present() -> None:
    """Test that import is not duplicated when already present."""
    code = """from __future__ import annotations
from contextvars import ContextVar

_context: ContextVar[SkyvernContext | None] = ContextVar("context")
"""
    module = cst.parse_module(code)
    result = add_future_annotations_import(module)

    # Should return the same module object (identity check)
    assert result is module
    assert result.code == code


def test_add_future_annotations_import_skips_when_not_needed() -> None:
    """Test that import is not added when all names are defined."""
    code = """from contextvars import ContextVar

class MyContext:
    pass

_context: ContextVar[MyContext] = ContextVar("context")
"""
    module = cst.parse_module(code)
    result = add_future_annotations_import(module)

    # Should return the same module object (identity check)
    assert result is module
    assert result.code == code


def test_add_future_annotations_import_skyvern_context_case() -> None:
    """Test the specific SkyvernContext case that was causing NameError."""
    code = """from contextvars import ContextVar

_context: ContextVar[SkyvernContext | None] = ContextVar(
    "Global context",
    default=None,
)


def current() -> SkyvernContext | None:
    return _context.get()
"""
    module = cst.parse_module(code)
    result = add_future_annotations_import(module)

    expected = """from __future__ import annotations
from contextvars import ContextVar

_context: ContextVar[SkyvernContext | None] = ContextVar(
    "Global context",
    default=None,
)


def current() -> SkyvernContext | None:
    return _context.get()
"""
    assert result.code == expected
    # Verify the result is valid Python
    cst.parse_module(result.code)
