"""Tests for dataclass constructor signature extraction."""

import libcst as cst

from aiservice.common.cst_utils import get_base_class_name, has_decorator
from aiservice.common.markdown_utils import extract_all_code_from_markdown
from core.languages.python.testgen.preprocessing.dataclass_constructor_notes import (
    FieldInfo,
    extract_dataclass_fields,
    find_all_dataclasses,
    format_constructor_signature,
    get_all_fields_with_inheritance,
    get_dataclass_constructor_notes,
)


class TestHasDecorator:
    def test_simple_dataclass_decorator(self) -> None:
        code = """
@dataclass
class Foo:
    x: int
"""
        tree = cst.parse_module(code)
        class_node = tree.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert has_decorator(class_node, "dataclass")

    def test_dataclass_with_call(self) -> None:
        code = """
@dataclass(frozen=True)
class Foo:
    x: int
"""
        tree = cst.parse_module(code)
        class_node = tree.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert has_decorator(class_node, "dataclass")

    def test_not_a_dataclass(self) -> None:
        code = """
class Foo:
    x: int
"""
        tree = cst.parse_module(code)
        class_node = tree.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert not has_decorator(class_node, "dataclass")

    def test_other_decorator(self) -> None:
        code = """
@other_decorator
class Foo:
    x: int
"""
        tree = cst.parse_module(code)
        class_node = tree.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert not has_decorator(class_node, "dataclass")


class TestGetBaseClassName:
    def test_simple_name(self) -> None:
        code = "class Foo(Bar): pass"
        tree = cst.parse_module(code)
        class_node = tree.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert get_base_class_name(class_node.bases[0]) == "Bar"

    def test_attribute(self) -> None:
        code = "class Foo(module.Bar): pass"
        tree = cst.parse_module(code)
        class_node = tree.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert get_base_class_name(class_node.bases[0]) == "Bar"


class TestExtractDataclassFields:
    def test_simple_fields(self) -> None:
        code = """
@dataclass
class Foo:
    x: int
    y: str
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        class_node = dataclasses["Foo"]
        fields = extract_dataclass_fields(class_node, source_lines, wrapper)

        assert len(fields) == 2
        assert fields[0].name == "x"
        assert fields[0].annotation == "int"
        assert not fields[0].has_default
        assert fields[1].name == "y"
        assert fields[1].annotation == "str"
        assert not fields[1].has_default

    def test_fields_with_defaults(self) -> None:
        code = """
@dataclass
class Foo:
    x: int
    y: str = "default"
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        class_node = dataclasses["Foo"]
        fields = extract_dataclass_fields(class_node, source_lines, wrapper)

        assert len(fields) == 2
        assert fields[0].name == "x"
        assert not fields[0].has_default
        assert fields[1].name == "y"
        assert fields[1].has_default
        assert fields[1].default_repr == '"default"'

    def test_complex_annotation(self) -> None:
        code = """
@dataclass
class Foo:
    items: list[str]
    mapping: dict[str, int] | None = None
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        class_node = dataclasses["Foo"]
        fields = extract_dataclass_fields(class_node, source_lines, wrapper)

        assert len(fields) == 2
        assert fields[0].name == "items"
        assert fields[0].annotation == "list[str]"
        assert fields[1].name == "mapping"
        assert "dict[str, int]" in fields[1].annotation


class TestFindAllDataclasses:
    def test_finds_single_dataclass(self) -> None:
        code = """
from dataclasses import dataclass

@dataclass
class Foo:
    x: int
"""
        result, source_lines, wrapper = find_all_dataclasses(code)
        assert "Foo" in result
        assert wrapper is not None
        assert len(source_lines) > 0

    def test_finds_multiple_dataclasses(self) -> None:
        code = """
from dataclasses import dataclass

@dataclass
class Foo:
    x: int

@dataclass
class Bar:
    y: str
"""
        result, _source_lines, wrapper = find_all_dataclasses(code)
        assert "Foo" in result
        assert "Bar" in result
        assert wrapper is not None

    def test_ignores_non_dataclasses(self) -> None:
        code = """
class NotADataclass:
    x: int

@dataclass
class IsADataclass:
    y: str
"""
        result, _source_lines, wrapper = find_all_dataclasses(code)
        assert "NotADataclass" not in result
        assert "IsADataclass" in result
        assert wrapper is not None

    def test_handles_syntax_error(self) -> None:
        code = "this is not valid python code {{{{"
        result, _source_lines, wrapper = find_all_dataclasses(code)
        assert result == {}
        assert wrapper is None


class TestGetAllFieldsWithInheritance:
    def test_no_inheritance(self) -> None:
        code = """
@dataclass
class Foo:
    x: int
    y: str
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        fields = get_all_fields_with_inheritance("Foo", dataclasses, source_lines, wrapper)

        assert len(fields) == 2
        assert all(f.source == "own" for f in fields)

    def test_simple_inheritance(self) -> None:
        code = """
@dataclass
class Base:
    a: int
    b: str

@dataclass
class Child(Base):
    c: float
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        fields = get_all_fields_with_inheritance("Child", dataclasses, source_lines, wrapper)

        assert len(fields) == 3
        assert fields[0].name == "a"
        assert fields[0].source == "inherited"
        assert fields[1].name == "b"
        assert fields[1].source == "inherited"
        assert fields[2].name == "c"
        assert fields[2].source == "own"

    def test_multi_level_inheritance(self) -> None:
        code = """
@dataclass
class GrandParent:
    a: int

@dataclass
class Parent(GrandParent):
    b: str

@dataclass
class Child(Parent):
    c: float
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        fields = get_all_fields_with_inheritance("Child", dataclasses, source_lines, wrapper)

        assert len(fields) == 3
        field_names = [f.name for f in fields]
        assert "a" in field_names
        assert "b" in field_names
        assert "c" in field_names

    def test_unknown_class(self) -> None:
        code = """
@dataclass
class Foo:
    x: int
"""
        dataclasses, source_lines, wrapper = find_all_dataclasses(code)
        assert wrapper is not None
        fields = get_all_fields_with_inheritance("Unknown", dataclasses, source_lines, wrapper)
        assert fields == []


class TestFormatConstructorSignature:
    def test_required_only(self) -> None:
        fields = [
            FieldInfo(name="x", annotation="int", has_default=False, default_repr=None, source="own"),
            FieldInfo(name="y", annotation="str", has_default=False, default_repr=None, source="own"),
        ]
        result = format_constructor_signature("Foo", fields)
        assert "Constructor signature for Foo:" in result
        assert "Required (positional) arguments:" in result
        assert "- x: int" in result
        assert "- y: str" in result
        assert "Optional" not in result

    def test_optional_only(self) -> None:
        fields = [
            FieldInfo(name="x", annotation="int", has_default=True, default_repr="0", source="own"),
            FieldInfo(name="y", annotation="str", has_default=True, default_repr='"default"', source="own"),
        ]
        result = format_constructor_signature("Foo", fields)
        assert "Constructor signature for Foo:" in result
        assert "Optional (keyword) arguments:" in result
        assert "- x: int = 0" in result
        assert '- y: str = "default"' in result
        assert "Required" not in result

    def test_mixed_fields(self) -> None:
        fields = [
            FieldInfo(name="a", annotation="int", has_default=False, default_repr=None, source="inherited"),
            FieldInfo(name="b", annotation="str", has_default=False, default_repr=None, source="own"),
            FieldInfo(name="c", annotation="float", has_default=True, default_repr="0.0", source="own"),
        ]
        result = format_constructor_signature("Foo", fields)
        assert "Required (positional) arguments:" in result
        assert "Optional (keyword) arguments:" in result
        assert "(from parent class)" in result

    def test_no_fields(self) -> None:
        result = format_constructor_signature("Empty", [])
        assert "no fields" in result


class TestExtractCodeFromMarkdown:
    def test_simple_code_block(self) -> None:
        markdown = """
Some text
```python
def foo():
    pass
```
More text
"""
        result = extract_all_code_from_markdown(markdown)
        assert "def foo():" in result
        assert "pass" in result
        assert "Some text" not in result

    def test_code_block_with_filepath(self) -> None:
        markdown = """
```python:path/to/file.py
class Foo:
    x: int
```
"""
        result = extract_all_code_from_markdown(markdown)
        assert "class Foo:" in result
        assert "x: int" in result

    def test_multiple_code_blocks(self) -> None:
        markdown = """
```python:file1.py
@dataclass
class A:
    x: int
```

```python:file2.py
@dataclass
class B:
    y: str
```
"""
        result = extract_all_code_from_markdown(markdown)
        assert "class A:" in result
        assert "class B:" in result


class TestGetDataclassConstructorNotes:
    def test_simple_dataclass(self) -> None:
        context = """
```python:models.py
from dataclasses import dataclass

@dataclass
class Config:
    name: str
    value: int
```
"""
        notes = get_dataclass_constructor_notes(context)
        assert len(notes) == 1
        assert "Constructor signature for Config:" in notes[0]
        assert "- name: str" in notes[0]
        assert "- value: int" in notes[0]

    def test_dataclass_with_inheritance(self) -> None:
        context = """
```python:models.py
from dataclasses import dataclass

@dataclass
class BaseConfig:
    model_name: str
    required_env_vars: list[str]

@dataclass
class ExtendedConfig(BaseConfig):
    extra_param: int = 0
```
"""
        notes = get_dataclass_constructor_notes(context)
        assert len(notes) == 2

        # Find the note for ExtendedConfig
        extended_note = next(n for n in notes if "ExtendedConfig" in n)
        assert "- model_name: str (from parent class)" in extended_note
        assert "- required_env_vars: list[str] (from parent class)" in extended_note
        assert "- extra_param: int = 0" in extended_note

    def test_no_dataclasses(self) -> None:
        context = """
def regular_function():
    pass
"""
        notes = get_dataclass_constructor_notes(context)
        assert notes == []

    def test_plain_code_without_markdown(self) -> None:
        context = """
from dataclasses import dataclass

@dataclass
class SimpleConfig:
    name: str
"""
        notes = get_dataclass_constructor_notes(context)
        assert len(notes) == 1
        assert "SimpleConfig" in notes[0]

    def test_llm_config_like_structure(self) -> None:
        """Test with a structure similar to the skyvern LLMConfig."""
        context = """
```python:skyvern/forge/sdk/api/llm/models.py
from dataclasses import dataclass, field
from typing import Optional

@dataclass(frozen=True)
class LLMConfigBase:
    model_name: str
    required_env_vars: list[str]
    supports_vision: bool
    add_assistant_prefix: bool

@dataclass(frozen=True)
class LLMConfig(LLMConfigBase):
    litellm_params: Optional[dict] = field(default=None)
    max_tokens: int | None = 4096
```
"""
        notes = get_dataclass_constructor_notes(context)

        # Should have notes for both classes
        assert len(notes) == 2

        # Find the LLMConfig note
        llm_config_note = next(n for n in notes if "LLMConfig:" in n)

        # Should show inherited fields from LLMConfigBase
        assert "- model_name: str (from parent class)" in llm_config_note
        assert "- required_env_vars: list[str] (from parent class)" in llm_config_note
        assert "- supports_vision: bool (from parent class)" in llm_config_note
        assert "- add_assistant_prefix: bool (from parent class)" in llm_config_note

        # Should also show own fields
        assert "- litellm_params:" in llm_config_note
        assert "- max_tokens:" in llm_config_note

    def test_full_output_for_complex_inheritance(self) -> None:
        """Test full generated notes output for complex dataclass inheritance."""
        context = """
```python:skyvern/forge/sdk/api/llm/models.py
from dataclasses import dataclass, field
from typing import Optional

@dataclass(frozen=True)
class LLMConfigBase:
    model_name: str
    required_env_vars: list[str]
    supports_vision: bool
    add_assistant_prefix: bool

@dataclass(frozen=True)
class LLMConfig(LLMConfigBase):
    litellm_params: Optional[dict] = field(default=None)
    max_tokens: int | None = 4096
    max_completion_tokens: int | None = None
    temperature: float | None = 0.7
    reasoning_effort: str | None = None

@dataclass(frozen=True)
class LLMRouterConfig(LLMConfigBase):
    model_list: list
    main_model_group: str
    redis_host: str | None = None
    redis_port: int | None = None
    fallback_model_group: str | None = None
    routing_strategy: str = "usage-based-routing"
    num_retries: int = 1
```
"""
        notes = get_dataclass_constructor_notes(context)

        # Should have notes for all 3 classes
        assert len(notes) == 3

        # Verify LLMConfigBase note (no inheritance)
        base_note = next(n for n in notes if "LLMConfigBase:" in n)
        expected_base = """Constructor signature for LLMConfigBase:
  Required (positional) arguments:
    - model_name: str
    - required_env_vars: list[str]
    - supports_vision: bool
    - add_assistant_prefix: bool"""
        assert base_note == expected_base

        # Verify LLMConfig note (with inheritance)
        config_note = next(n for n in notes if "LLMConfig:" in n)
        expected_config = """Constructor signature for LLMConfig:
  Required (positional) arguments:
    - model_name: str (from parent class)
    - required_env_vars: list[str] (from parent class)
    - supports_vision: bool (from parent class)
    - add_assistant_prefix: bool (from parent class)
  Optional (keyword) arguments:
    - litellm_params: Optional[dict] = field(default=None)
    - max_tokens: int | None = 4096
    - max_completion_tokens: int | None = None
    - temperature: float | None = 0.7
    - reasoning_effort: str | None = None"""
        assert config_note == expected_config

        # Verify LLMRouterConfig note (with inheritance + own required fields)
        router_note = next(n for n in notes if "LLMRouterConfig:" in n)
        expected_router = """Constructor signature for LLMRouterConfig:
  Required (positional) arguments:
    - model_name: str (from parent class)
    - required_env_vars: list[str] (from parent class)
    - supports_vision: bool (from parent class)
    - add_assistant_prefix: bool (from parent class)
    - model_list: list
    - main_model_group: str
  Optional (keyword) arguments:
    - redis_host: str | None = None
    - redis_port: int | None = None
    - fallback_model_group: str | None = None
    - routing_strategy: str = "usage-based-routing"
    - num_retries: int = 1"""
        assert router_note == expected_router
