# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

import json
import logging
from pathlib import Path

from autogen.import_utils import optional_import_block

with optional_import_block() as result:
    from fastapi_code_generator.parser import OpenAPIParser, Operation
    from fastapi_code_generator.visitor import Visitor

from pydantic import BaseModel

from autogen.agentchat.conversable_agent import ConversableAgent
from autogen.llm_config import LLMConfig


class Group(BaseModel):
    name: str
    description: str


class GroupSuggestions(BaseModel):
    groups: list[Group]


class GroupNames(BaseModel):
    groups: list[str]


logger = logging.getLogger(__name__)

GROUP_DISCOVERY_MESSAGE = (
    "You are a senior Python engineer. You will be shown a batch of API functions. "
    "These functions are not guaranteed to be related — your task is to analyze them individually and find meaningful groups *within the batch*.\n\n"
    "You should propose a list of group names, and for each group, provide a short description of the kind of functions it contains.\n\n"
    "How to group:\n"
    "- Focus on what the function operates on — e.g., functions that manipulate a board go in 'board_operations', functions related to users in 'user_management', etc.\n"
    "- Do NOT assume the entire batch forms a single group.\n"
    "- Do NOT use generic categories like 'misc', 'utils', or 'general'.\n"
    "- Favor *granular but meaningful* groups. For instance, 'user_profile_handling' and 'user_authentication' could be separate if their purposes are distinct.\n\n"
    "Formatting:\n"
    "- Group names must be short, descriptive, and in snake_case.\n"
    "- Return a list of group names and a 1-2 sentence description of what each group includes.\n"
)


GROUP_ASSIGNMENT_MESSAGE = (
    "You are a senior Python engineer. You will be given a function description "
    "and a list of possible groups. Choose the most suitable groups for the function.\n"
    "- A function can belong to multiple groups.\n"
    "- The groups should be relevant to the function's purpose and functionality.\n"
    "- You MUST choose a group from the possible groups, you cannot respond with empty grouping any equivalent.\n"
)


def chunk_list(items: list, size: int) -> list[list]:
    return [items[i : i + size] for i in range(0, len(items), size)]


def discover_groups(operations: list["Operation"], chunk_size: int = 30) -> dict[str, str]:
    llm_config = LLMConfig.get_current_llm_config().copy()

    for config in llm_config.config_list:
        config.response_format = GroupSuggestions

    agent = ConversableAgent(
        name="group_discovery_agent",
        system_message=GROUP_DISCOVERY_MESSAGE,
        llm_config=llm_config,
    )
    groups = {}

    for chunk in chunk_list(operations, chunk_size):
        func_descriptions = [f"- {op.function_name}: {op.summary} (args: {op.arguments})" for op in chunk]
        message = "Here are some functions:\n" + "\n".join(func_descriptions)

        response = agent.run(message=message, max_turns=1, user_input=False)

        for event in response.events:
            if event.type == "text" and event.content.sender == "group_discovery_agent":
                # Naively parse "group_name: description" from text block
                new_groups = GroupSuggestions.model_validate_json(event.content.content).groups
                groups.update(new_groups)

    logger.warning("Discovered groups: %s", groups)

    # Remove duplicates
    agent = ConversableAgent(
        name="group_refining_agent",
        system_message=GROUP_DISCOVERY_MESSAGE,
        llm_config=llm_config,
    )

    message = (
        "You need to refine the group names and descriptions to ensure they are unique.\n"
        "Here are the groups:\n" + "\n".join([f"- {name}: {desc}" for name, desc in groups.items()])
    )
    response = agent.run(message=message, max_turns=1, user_input=False)
    for event in response.events:
        if event.type == "text" and event.content.sender == "group_refining_agent":
            # Naively parse "group_name: description" from text block
            refined_groups = json.loads(event.content.content)

    return refined_groups


def assign_operation_to_group(operation: "Operation", groups: dict[str, str]) -> str:
    llm_config = LLMConfig.get_current_llm_config().copy()

    for config in llm_config.config_list:
        config.response_format = GroupNames

    agent = ConversableAgent(
        name="group_assignment_agent",
        system_message=GROUP_ASSIGNMENT_MESSAGE,
        llm_config=llm_config,
    )

    message = (
        "Function summary:\n"
        f"{operation.summary}\n\n"
        f"Arguments: {operation.arguments}\n\n"
        f"Available groups: {json.dumps(groups)}\n\n"
        "What group should this function go in?"
    )

    response = agent.run(message=message, max_turns=1, user_input=True)

    groups = []
    for event in response.events:
        if event.type == "text" and event.content.sender == "group_assignment_agent":
            groups = GroupNames.model_validate_json(event.content.content).groups

    return groups


def refine_group_names(groups: dict[str, str]) -> dict[str, str]:
    # Optional: normalize names, merge similar ones (e.g., using embeddings or string similarity)
    # Placeholder for now:
    return groups


def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> dict[str, object]:
    operations = sorted(parser.operations.values(), key=lambda op: op.path)

    # ---- PASS 1: DISCOVER GROUPS ----
    logger.warning("Discovering groups...")
    discovered_groups = discover_groups(operations)
    logger.warning("Discovered groups: %s", discovered_groups)

    # ---- PASS 2: ASSIGN OPERATIONS TO GROUPS ----
    logger.warning("Assigning operations to groups...")
    for op in operations:
        logger.warning("Assigning operation %s to groups...", op.function_name)
        groups = assign_operation_to_group(op, discovered_groups)
        op.tags = groups
        logger.warning("Assigned groups: %s", groups)

    return {"operations": operations}


visit: "Visitor" = custom_visitor
