# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
#
# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
# SPDX-License-Identifier: MIT
import io
import json
import logging
import os
import sqlite3
import uuid
from collections.abc import Callable
from pathlib import Path
from typing import Any
from unittest.mock import Mock, patch

import pytest

from autogen.import_utils import optional_import_block, run_for_optional_imports

with optional_import_block() as result:
    import openai  # noqa: F401
    from openai import AzureOpenAI


import autogen.runtime_logging
from autogen.logger.logger_utils import EventStreamHandler, event_print, get_current_ts, to_dict

SAMPLE_CHAT_REQUEST = json.loads(
    """
{
    "messages": [
        {
            "content": "You are roleplaying a high school student strugling with linear algebra. Regardless how well the teacher explains things to you, you just don't quite get it. Keep your questions short.",
            "role": "system"
        },
        {
            "content": "Can you explain the difference between eigenvalues and singular values again?",
            "role": "assistant"
        },
        {
            "content": "Certainly!\\n\\nEigenvalues are associated with square matrices. They are the scalars, \\u03bb, that satisfy the equation\\n\\nA*x = \\u03bb*x\\n\\nwhere A is a square matrix, x is a nonzero vector (the eigenvector), and \\u03bb is the eigenvalue. The eigenvalue equation shows how the vector x is stretched or shrunk by the matrix A.\\n\\nSingular values, on the other hand, are associated with any m x n matrix, whether square or rectangular. They come from the matrix's singular value decomposition (SVD) and are the square roots of the non-negative eigenvalues of the matrix A*A^T or A^T*A (where A^T is the transpose of A). Singular values, denoted often by \\u03c3, represent the magnitude of the principal axes of the data's distribution and are always non-negative.\\n\\nTo sum up, eigenvalues relate to how a matrix scales vectors (specific to square matrices), while singular values give a measure of how a matrix stretches space (applicable to all matrices).",
            "role": "user"
        }
    ],
    "model": "gpt-4o"
}
"""
)

SAMPLE_CHAT_RESPONSE = json.loads(
    """
{
    "id": "chatcmpl-8k57oSg1fz2JwpMcEOWMqUvwjf0cb",
    "choices": [
        {
            "finish_reason": "stop",
            "index": 0,
            "logprobs": null,
            "message": {
                "content": "Oh, wait, I don't think I completely understand the concept of matrix multiplication. Could you break down how you multiply two matrices together?",
                "role": "assistant",
                "function_call": null,
                "tool_calls": null
            }
        }
    ],
    "created": 1705993480,
    "model": "gpt-4o",
    "object": "chat.completion",
    "system_fingerprint": "fp_6d044fb900",
    "usage": {
        "completion_tokens": 28,
        "prompt_tokens": 274,
        "total_tokens": 302
    }
}
"""
)


def dummy_function(param1: str, param2: int) -> Any:
    return param1 * param2


###############################################################


@pytest.fixture(scope="function")
def db_connection():
    autogen.runtime_logging.start(config={"dbname": ":memory:"})
    con = autogen.runtime_logging.get_connection()
    con.row_factory = sqlite3.Row
    yield con

    autogen.runtime_logging.stop()


def get_sample_chat_completion(response):
    return {
        "invocation_id": str(uuid.uuid4()),
        "client_id": 140609438577184,
        "wrapper_id": 140610167717744,
        "request": SAMPLE_CHAT_REQUEST,
        "response": response,
        "is_cached": 0,
        "cost": 0.347,
        "start_time": get_current_ts(),
        "agent": autogen.AssistantAgent(name="TestAgent", code_execution_config=False),
    }


@pytest.mark.parametrize(
    "response, expected_logged_response",
    [
        (SAMPLE_CHAT_RESPONSE, SAMPLE_CHAT_RESPONSE),
        (None, {"response": None}),
        ("error in response", {"response": "error in response"}),
    ],
)
def test_log_completion(response, expected_logged_response, db_connection):
    cur = db_connection.cursor()

    sample_completion = get_sample_chat_completion(response)
    autogen.runtime_logging.log_chat_completion(**sample_completion)

    query = """
        SELECT invocation_id, client_id, wrapper_id, request, response, is_cached,
            cost, start_time, source_name FROM chat_completions
    """

    for row in cur.execute(query):
        assert row["invocation_id"] == sample_completion["invocation_id"]
        assert row["client_id"] == sample_completion["client_id"]
        assert row["wrapper_id"] == sample_completion["wrapper_id"]
        assert json.loads(row["request"]) == sample_completion["request"]
        assert json.loads(row["response"]) == expected_logged_response
        assert row["is_cached"] == sample_completion["is_cached"]
        assert row["cost"] == sample_completion["cost"]
        assert row["start_time"] == sample_completion["start_time"]
        assert row["source_name"] == "TestAgent"


def test_log_function_use(db_connection):
    cur = db_connection.cursor()

    source = autogen.AssistantAgent(name="TestAgent", code_execution_config=False)
    func: Callable[[str, int], Any] = dummy_function
    args = {"foo": "bar"}
    returns = True

    autogen.runtime_logging.log_function_use(agent=source, function=func, args=args, returns=returns)

    query = """
        SELECT source_id, source_name, function_name, args, returns, timestamp
        FROM function_calls
    """

    for row in cur.execute(query):
        assert row["source_name"] == "TestAgent"
        assert row["args"] == json.dumps(args)
        assert row["returns"] == json.dumps(returns)


@run_for_optional_imports(["openai"], "openai")
def test_log_new_agent(db_connection):
    from autogen import AssistantAgent

    cur = db_connection.cursor()
    agent_name = "some_assistant"
    config_list = [{"model": "gpt-4o", "api_key": "some_key"}]

    agent = AssistantAgent(agent_name, llm_config={"config_list": config_list})
    init_args = {"foo": "bar", "baz": {"other_key": "other_val"}, "a": None}

    autogen.runtime_logging.log_new_agent(agent, init_args)

    query = """
        SELECT session_id, name, class, init_args FROM agents
    """

    for row in cur.execute(query):
        assert row["session_id"] and str(uuid.UUID(row["session_id"], version=4)) == row["session_id"], (
            "session id is not valid uuid"
        )
        assert row["name"] == agent_name
        assert row["class"] == "AssistantAgent"
        assert row["init_args"] == json.dumps(init_args)


@run_for_optional_imports(["openai"], "openai")
def test_log_oai_wrapper(db_connection):
    from autogen import OpenAIWrapper

    cur = db_connection.cursor()

    llm_config = {"config_list": [{"model": "gpt-4o", "api_key": "some_key", "base_url": "some url"}]}
    init_args = {"llm_config": llm_config, "base_config": {}}
    wrapper = OpenAIWrapper(**llm_config)

    autogen.runtime_logging.log_new_wrapper(wrapper, init_args)

    query = """
        SELECT session_id, init_args FROM oai_wrappers
    """

    for row in cur.execute(query):
        assert row["session_id"] and str(uuid.UUID(row["session_id"], version=4)) == row["session_id"], (
            "session id is not valid uuid"
        )
        saved_init_args = json.loads(row["init_args"])
        assert "config_list" in saved_init_args
        assert "api_key" not in saved_init_args["config_list"][0]
        assert "base_url" not in saved_init_args["config_list"][0]
        assert "base_config" in saved_init_args


@run_for_optional_imports(["openai"], "openai")
def test_log_oai_client(db_connection):
    cur = db_connection.cursor()

    openai_config = {
        "api_key": "some_key",
        "api_version": "2024-08-06",
        "azure_deployment": "gpt-4o",
        "azure_endpoint": "https://foobar.openai.azure.com/",
    }
    client = AzureOpenAI(**openai_config)

    autogen.runtime_logging.log_new_client(client, Mock(), openai_config)

    query = """
        SELECT session_id, init_args, class FROM oai_clients
    """

    for row in cur.execute(query):
        assert row["session_id"] and str(uuid.UUID(row["session_id"], version=4)) == row["session_id"], (
            "session id is not valid uuid"
        )
        assert row["class"] == "AzureOpenAI"
        saved_init_args = json.loads(row["init_args"])
        assert "api_version" in saved_init_args
        assert "api_key" not in saved_init_args


def test_to_dict():
    from autogen import Agent
    from autogen.coding import LocalCommandLineCodeExecutor

    agent_executor = LocalCommandLineCodeExecutor(work_dir=Path())

    agent1 = autogen.ConversableAgent(
        "alice",
        human_input_mode="NEVER",
        llm_config=False,
        default_auto_reply="This is alice speaking.",
        code_execution_config={"executor": agent_executor},
    )

    agent2 = autogen.ConversableAgent(
        "bob",
        human_input_mode="NEVER",
        llm_config=False,
        default_auto_reply="This is bob speaking.",
        function_map={"test_func": lambda x: x},
    )

    class Foo:
        def __init__(self):
            self.a = 1.234
            self.b = "some string"
            self.c = {"some_key": [7, 8, 9]}
            self.d = None
            self.test_function = lambda x, y: x + y
            self.extra_key = "remove this key"
            self.path = Path("/to/something")

    class Bar:
        def init(self):
            pass

        def build(self):
            self.foo_val = [Foo()]
            self.o = {"key_1": None, "key_2": [{"nested_key_1": ["nested_val_1", "nested_val_2"]}]}
            self.agents = [agent1, agent2]
            self.first_agent = agent1

    bar = Bar()
    bar.build()

    expected_path = "\\to\\something" if os.name == "nt" else "/to/something"
    expected_foo_val_field = [
        {
            "a": 1.234,
            "b": "some string",
            "c": {"some_key": [7, 8, 9]},
            "d": None,
            "test_function": "self.test_function = lambda x, y: x + y",
            "path": expected_path,
        }
    ]

    expected_o_field = {"key_2": [{"nested_key_1": ["nested_val_1", "nested_val_2"]}]}

    result = to_dict(bar, exclude=("key_1", "extra_key"), no_recursive=(Agent))
    assert result["foo_val"] == expected_foo_val_field
    assert result["o"] == expected_o_field
    assert len(result["agents"]) == 2
    assert result["agents"][0] == "alice"
    assert result["agents"][1] == "bob"
    assert "alice" in result["first_agent"]


@patch("logging.Logger.error")
def test_logging_exception_will_not_crash_only_print_error(mock_logger_error, db_connection):
    sample_completion = get_sample_chat_completion(SAMPLE_CHAT_REQUEST)
    sample_completion["is_cached"] = {"foo": "bar"}

    autogen.runtime_logging.log_chat_completion(**sample_completion)

    args, _ = mock_logger_error.call_args
    error_message = args[0]
    assert error_message.startswith("[sqlite logger]Error running query with query")


class _RecordingHandler(logging.Handler):
    def __init__(self) -> None:
        super().__init__()
        self.records: list[logging.LogRecord] = []

    def emit(self, record: logging.LogRecord) -> None:
        self.records.append(record)


def test_event_print_with_custom_logger_and_handler() -> None:
    handler = _RecordingHandler()
    logger = logging.getLogger("ag2.event.processor.test")
    logger.handlers = [handler]
    logger.propagate = False

    event_print("hello", "world", sep="-", end="!", flush=False, logger=logger, level=logging.WARNING)

    assert len(handler.records) == 1
    record = handler.records[0]
    assert record.getMessage() == "hello-world"
    assert getattr(record, "ag2_event_end") == "!"
    assert getattr(record, "ag2_event_flush") is False
    assert record.levelno == logging.WARNING


def test_event_print_default_logger_respects_end_and_flush() -> None:
    stream = io.StringIO()
    logger = logging.getLogger("ag2.event.processor")
    old_handlers = logger.handlers[:]
    old_propagate = logger.propagate
    try:
        for handler in logger.handlers[:]:
            logger.removeHandler(handler)
        handler = EventStreamHandler(stream)
        handler.setFormatter(logging.Formatter("%(message)s"))
        logger.addHandler(handler)
        logger.setLevel(logging.INFO)
        logger.propagate = False

        event_print("structured", "output", sep="|", end="END", flush=True)

        assert stream.getvalue() == "structured|outputEND"
    finally:
        for handler in logger.handlers[:]:
            logger.removeHandler(handler)
        for handler in old_handlers:
            logger.addHandler(handler)
        logger.propagate = old_propagate
