# Copyright 2025 TOYOTA MOTOR CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

from typing import Union, List, Tuple, Dict, Optional, Iterator, Any, Generic, TypeVar, get_origin
import sys
import os
import multiprocessing
import time
import math
import gc
import weakref
import pytest
import egrpc

env_name = "client"

def serve(port: int, error_queue: Any) -> None:
    global env_name
    env_name = "server"

    # redirect stdout to parent process
    sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)

    try:
        egrpc.serve(port)
    except AssertionError as e:
        error_queue.put(str(e))
    except Exception as e:
        error_queue.put(f"Unexpected error: {e}")

@pytest.fixture(scope="module")
def server() -> Iterator[None]:
    port = 12345

    error_queue: Any = multiprocessing.Queue()

    server_process = multiprocessing.Process(target=serve, args=(port, error_queue))
    server_process.start()

    time.sleep(1)

    if not server_process.is_alive():
        raise RuntimeError("Server failed to start.")

    egrpc.connect("localhost", port)

    try:
        yield

    finally:
        while not error_queue.empty():
            error_message = error_queue.get()
            pytest.fail(f"Server process error: {error_message}")

        server_process.terminate()
        server_process.join()

        egrpc.disconnect()

@egrpc.function
def get_gvar() -> str:
    return env_name

def test_remote_exec(server: Any) -> None:
    assert env_name == "client"
    assert get_gvar() == "server"

@egrpc.function
def func1(name: str, age: int) -> str:
    return f"{name}: {age}"

@egrpc.function
def func2(name: str, weight: Union[int, float]) -> str:
    return f"{name}: {weight:.2f}"

@egrpc.function
def func3(x: Union[int, float]) -> Union[int, float]:
    return x * x

@egrpc.function
def func4(lst: Union[List[int], List[str]], n: int) -> Union[List[int], List[str]]:
    return lst * n

@egrpc.function
def func5(lst: List[Union[int, str]], n: int) -> List[Union[int, str]]:
    return lst * n

@egrpc.function
def func6(tup: Tuple[int, str]) -> Tuple[str, int]:
    return tup[1], tup[0]

@egrpc.function
def func7(d: Dict[str, int]) -> Dict[str, List[str]]:
    return {k: [k] * v for k, v in d.items()}

@egrpc.function
def func8(x: Optional[int] = None) -> Optional[int]:
    return x * x if x is not None else None

@egrpc.function
def func9(x: Tuple[int, ...]) -> Tuple[float, ...]:
    return tuple(float(n) for n in x)

@egrpc.function
def func_slice(s: slice) -> slice:
    return s

EllipsisType = type(...)

@egrpc.function
def func_ellipsis(e: EllipsisType) -> EllipsisType:  # type: ignore[valid-type]
    return e

@egrpc.function
def func_slice_tuple(idx: Tuple[Union[int, slice], ...]) -> Tuple[Union[int, slice], ...]:
    return idx

def test_function(server: Any) -> None:
    assert func1("Alice", 30) == "Alice: 30"
    assert func2("Bob", 60) == "Bob: 60.00"
    assert func2("Bob", 60.2) == "Bob: 60.20"
    assert func3(2) == 4
    assert func3(4.3) == pytest.approx(4.3 * 4.3)
    assert func4([1, 2, 3], 2) == [1, 2, 3, 1, 2, 3]
    assert func4(["a", "b", "c"], 2) == ["a", "b", "c", "a", "b", "c"]
    assert func5([1, "b", 3], 2) == [1, "b", 3, 1, "b", 3]
    assert func6((1, "a")) == ("a", 1)
    assert func7({"a": 1, "b": 2}) == {"a": ["a"], "b": ["b", "b"]}
    assert func8(2) == 4
    assert func8() is None
    assert func9((1, 2, 3)) == (1.0, 2.0, 3.0)
    assert func_slice(slice(1, 10, 2)) == slice(1, 10, 2)
    assert func_slice(slice(None, 5)) == slice(None, 5)
    assert func_slice(slice(3, None)) == slice(3, None)
    assert func_slice(slice(None)) == slice(None)
    assert func_ellipsis(...) is ...
    assert func_slice_tuple((slice(None), slice(None))) == (slice(None), slice(None))

@egrpc.dataclass
class Data():
    x: int
    y: float

@egrpc.function
def dcfunc1(d: Data) -> float:
    return d.x * d.y

@egrpc.function
def dcfunc2(d1: Data, d2: Data) -> bool:
    return d1 == d2

def test_dataclass(server: Any) -> None:
    d = Data(3, 2.2)
    assert dcfunc1(d) == pytest.approx(3 * 2.2)
    assert dcfunc2(d, d) is True
    assert dcfunc2(d, Data(2, 2.2)) is False

@egrpc.remoteclass
class Point():
    @egrpc.method
    def __init__(self, x: Union[int, float], y: Union[int, float], z: Union[int, float]):
        self.x = x
        self.y = y
        self.z = z

    @egrpc.method
    def norm(self) -> float:
        return math.sqrt(self.x * self.x + self.y * self.y + self.z * self.z)

    @egrpc.method
    def __str__(self) -> str:
        return ",".join([str(self.x), str(self.y), str(self.z)])

@egrpc.function
def identity(p: Point) -> Point:
    return p

@egrpc.function
def norm(p: Point) -> float:
    return p.norm()

def test_remoteclass(server: Any) -> None:
    p1 = Point(1, 2, 3)
    p2 = Point(1.1, 2.2, 3.3)

    assert p1 != p2
    assert identity(p1) == p1
    assert p1.norm() == pytest.approx(math.sqrt(14))
    assert norm(p1) == pytest.approx(math.sqrt(14))
    assert str(p1) == "1,2,3"

    with pytest.raises(AttributeError):
        p1.x

@egrpc.remoteclass
class Hoge():
    def __init__(self, x: int):
        self.x = x
        self._fuga = Fuga(str(x))

    @egrpc.property
    def value(self) -> int:
        return self.x

    @value.setter
    def value(self, x: int) -> None:
        self.x = x
        self._fuga.value = str(x)

    @egrpc.method
    def fuga(self) -> "Fuga":
        return self._fuga

@egrpc.remoteclass
class Fuga():
    def __init__(self, x: str):
        self.x = x

    @egrpc.property
    def value(self) -> str:
        return self.x

    @value.setter
    def value(self, x: str) -> None:
        self.x = x

    @egrpc.method
    def hoge(self) -> "Hoge":
        return Hoge(int(self.x))

@egrpc.function
def gen_hoge(x: int) -> Hoge:
    return Hoge(x)

def test_remoteclass2(server: Any) -> None:
    hoge = gen_hoge(1)
    fuga = hoge.fuga()
    hoge_ = fuga.hoge()
    assert hoge.value == hoge_.value

    hoge.value = 2
    assert hoge.value == 2
    assert hoge.fuga().value == "2"

def test_remoteclass3(server: Any) -> None:
    hoge = gen_hoge(7)
    fuga1 = hoge.fuga()
    fuga2 = hoge.fuga()
    assert fuga1 is fuga2
    assert fuga1.value == "7"

    ref = weakref.ref(fuga1)
    del fuga1
    del fuga2
    gc.collect()
    assert ref() is None

    fuga3 = hoge.fuga()
    assert fuga3.value == "7"

@egrpc.remoteclass
class Animal:
    def __init__(self, name: str):
        self._name = name

    @egrpc.property
    def name(self) -> str:
        return self._name

    @egrpc.method
    def speak(self) -> str:
        return "..."

@egrpc.remoteclass
class Dog(Animal):
    def __init__(self, name: str):
        super().__init__(name)

    @egrpc.method
    def speak(self) -> str:
        return "woof"

@egrpc.remoteclass
class Cat(Animal):
    def __init__(self, name: str):
        super().__init__(name)

    @egrpc.method
    def speak(self) -> str:
        return "meow"

@egrpc.remoteclass
class Mammal(Animal):
    def __init__(self, name: str, legs: int):
        super().__init__(name)
        self._legs = legs

    @egrpc.property
    def legs(self) -> int:
        return self._legs

@egrpc.remoteclass
class Horse(Mammal):
    def __init__(self, name: str):
        super().__init__(name, 4)

    @egrpc.method
    def speak(self) -> str:
        return "neigh"

@egrpc.function
def create_animal(kind: str, name: str) -> Animal:
    if kind == "dog":
        return Dog(name)
    elif kind == "cat":
        return Cat(name)
    elif kind == "horse":
        return Horse(name)
    else:
        return Animal(name)

@egrpc.function
def get_animal_name(animal: Animal) -> str:
    return animal.name

def test_inheritance(server: Any) -> None:
    dog = create_animal("dog", "Pochi")
    assert isinstance(dog, Dog)
    assert dog.name == "Pochi"
    assert dog.speak() == "woof"
    assert get_animal_name(dog) == "Pochi"

    cat = create_animal("cat", "Tama")
    assert isinstance(cat, Cat)
    assert cat.name == "Tama"
    assert cat.speak() == "meow"
    assert get_animal_name(cat) == "Tama"

    horse = create_animal("horse", "Uma")
    assert isinstance(horse, Horse)
    assert horse.name == "Uma"
    assert horse.speak() == "neigh"
    assert horse.legs == 4
    assert get_animal_name(horse) == "Uma"

    animal = create_animal("unknown", "X")
    assert isinstance(animal, Animal)
    assert animal.name == "X"
    assert animal.speak() == "..."

@egrpc.function
def raise_value_error(msg: str) -> None:
    raise ValueError(msg)

@egrpc.function
def raise_type_error(msg: str) -> None:
    raise TypeError(msg)

class CustomError(Exception):
    pass

@egrpc.function
def raise_custom_error(msg: str) -> None:
    raise CustomError(msg)

@egrpc.remoteclass
class ErrorThrower:
    @egrpc.method
    def __init__(self, value: int):
        self._value = value

    @egrpc.method
    def raise_if_negative(self) -> int:
        if self._value < 0:
            raise ValueError(f"value must be non-negative, got {self._value}")
        return self._value

def test_exception_propagation(server: Any) -> None:
    with pytest.raises(ValueError) as exc_info_ve:
        raise_value_error("test value error")
    assert str(exc_info_ve.value) == "test value error"

    with pytest.raises(TypeError) as exc_info_te:
        raise_type_error("test type error")
    assert str(exc_info_te.value) == "test type error"

    with pytest.raises(CustomError) as exc_info_ce:
        raise_custom_error("test custom error")
    assert str(exc_info_ce.value) == "test custom error"

def test_exception_propagation_method(server: Any) -> None:
    obj = ErrorThrower(10)
    assert obj.raise_if_negative() == 10

    obj_neg = ErrorThrower(-5)
    with pytest.raises(ValueError) as exc_info:
        obj_neg.raise_if_negative()
    assert "value must be non-negative" in str(exc_info.value)

_T = TypeVar("_T")

class _MyInt(int, Generic[_T]):
    pass

egrpc.register_primitive_type(_MyInt, "int64")

def _my_int_subtype_rule(t1: Any, t2: Any) -> Optional[bool]:
    o1 = get_origin(t1)
    o2 = get_origin(t2)
    if o1 is _MyInt and o2 is None:
        return t2 in (int, float)
    if o1 is None and o2 is _MyInt:
        return False
    return None

egrpc.register_subtype_rule(_my_int_subtype_rule)

@egrpc.function
def my_int_identity(x: _MyInt[Any]) -> _MyInt[Any]:
    return x

@egrpc.function
def accept_my_int_or_float(x: Union[_MyInt[Any], float]) -> float:
    return float(x)

def test_custom_type_registration(server: Any) -> None:
    assert my_int_identity(_MyInt(42)) == 42
    assert accept_my_int_or_float(_MyInt(7)) == 7.0
    assert accept_my_int_or_float(3.14) == pytest.approx(3.14)
