from typing import Literal, TypedDict

from openai.types.chat import ChatCompletionUserMessageParam
from pydantic import BaseModel, Field

from tfy_assistant_framework.assistant import (
    AssistantUpdateToUser,
    send_update,
)
from tfy_assistant_framework.swarm import (
    Agent,
    AgentResponseValidationError,
    ContextVarType,
    get_response_from_structured_agent,
)


# Tools and state setup
class AccountInfo(TypedDict):
    balance: float
    name: str


ACCOUNTS: dict[int, AccountInfo] = {
    1: {"balance": 1000.0, "name": "John"},
    2: {"balance": 500.0, "name": "Alice"},
}


def is_account_valid(account_id: int) -> bool:
    return account_id in ACCOUNTS


def get_balance(context_variables: ContextVarType) -> float:
    return ACCOUNTS[context_variables["initiator_account_id"]]["balance"]


def make_transfer(
    context_variables: ContextVarType, to_account: int, amount: float
) -> str:
    sender = ACCOUNTS[context_variables["initiator_account_id"]]
    recipient = ACCOUNTS.get(to_account)

    if not recipient:
        raise ValueError("Recipient not found")
    if amount > sender["balance"]:
        raise ValueError("Insufficient funds")

    sender["balance"] -= amount
    recipient["balance"] += amount
    return f"Transferred ${amount} to {recipient['name']}"


# Agents Setup
class TransferIntent(BaseModel):
    intent_type: Literal["transfer"] = Field(
        ..., description="Intent type, only accepts 'transfer'"
    )
    amount: float = Field(..., description="Transfer amount", gt=0)
    to_account: int = Field(..., description="Recipient account ID")


class ExitIntent(BaseModel):
    intent_type: Literal["exit"] = Field(
        ..., description="Intent type, only accepts 'exit'"
    )


class IntentAgentResponse(BaseModel):
    intent: TransferIntent | ExitIntent


class TransferResultResponse(BaseModel):
    message: str = Field(..., description="Description of the transfer")


class IntentAgent(Agent[IntentAgentResponse]):
    async def validate_response(
        self,
        response: IntentAgentResponse,
        context_variables: ContextVarType,
    ) -> AgentResponseValidationError | IntentAgentResponse:
        if response.intent.intent_type == "transfer":  # acts as a hard gaurd rail
            if not is_account_valid(response.intent.to_account):
                return AgentResponseValidationError(
                    error_message="Recipient account is not valid",
                )
            if (
                response.intent.amount
                > ACCOUNTS[context_variables["initiator_account_id"]]["balance"]
            ):
                return AgentResponseValidationError(
                    error_message="Insufficient funds",
                )

        return response


class TransferAgent(Agent[TransferResultResponse]): ...


intent_agent = IntentAgent(
    response_type=IntentAgentResponse,
    instructions="Interact with user and Identify banking intent from user request",
    functions=[get_balance, is_account_valid],
)

transfer_agent = TransferAgent(
    response_type=TransferResultResponse,
    instructions="Handle money transfers between accounts",
    functions=[make_transfer, get_balance],
)


class CustomerSupportRequest(BaseModel):
    account_id: int


# Run the customer support
async def run_customer_support_agent(body: CustomerSupportRequest) -> None:
    context_variables = {"initiator_account_id": body.account_id}

    response = await get_response_from_structured_agent(
        intent_agent, context_variables=context_variables, interactive=True
    )

    if response.intent.intent_type == "transfer":
        result = await get_response_from_structured_agent(
            transfer_agent,
            context_variables=context_variables,
            interactive=False,
            messages=[
                ChatCompletionUserMessageParam(
                    content=response.intent.model_dump_json(), role="user"
                )
            ],
        )
        await send_update(AssistantUpdateToUser(content=result.message))

    await send_update(AssistantUpdateToUser(content="Thank you for contacting us"))


async def main() -> None:
    await run_customer_support_agent(CustomerSupportRequest(account_id=1))


if __name__ == "__main__":
    import asyncio

    result = asyncio.run(main())
    print(result)
