from backend.smart_domain.api_layer_code_tool import apiLayerCodeGenerator
from backend.smart_domain.association import associationCodeGenerator
from backend.smart_domain.association_impl import asociationImplCodeGenerator
from backend.smart_domain.db_entity_repository import dbEntityRepositoryCodeGenerator
from backend.smart_domain.domain_layer_code_tool import domainLayerCodeGenerator
from backend.smart_domain.entity import entityCodeGenerator
from backend.smart_domain.persistent_layer_code_tool import persistentLayerCodeGenerator
from langgraph.prebuilt import create_react_agent

from backend.models import model, models
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.sqlite import SqliteSaver


memory = SqliteSaver.from_conn_string(":memory:")

code_agent_tools = [domainLayerCodeGenerator, entityCodeGenerator, associationCodeGenerator, persistentLayerCodeGenerator, dbEntityRepositoryCodeGenerator, asociationImplCodeGenerator, apiLayerCodeGenerator]

agent_executor = create_react_agent(model, code_agent_tools, checkpointer=memory)

async def code_generator(
    model,
    command: str,
    history,
    context: str):
    message = command + f"""
    context: '''{context}'''
"""

    history[-1][1] = ""
    providor = models[model] if model in models else "openai"
    async for event in agent_executor.astream_events(
        {"messages": [HumanMessage(content=message)]}, 
        {'configurable': {
            "thread_id": "1",
            "model": model,
            "model_provider": providor
            }
        }, version="v1"):
        kind = event["event"]
        if kind == "on_chat_model_stream":
            content = event["data"]["chunk"].content
            if content:
                history[-1][1] += content
                yield history
    yield history

if __name__ == '__main__':
    message = """generate entity code according the following context: 
    '''Feature Config指的是一些特性的配置信息，例如应用程序的功能模块和设置。

管理员可以通过界面进行新增、修改、分页查看、发布、撤销和删除操作；

新增Feature Config应该提供 FeatureKey（通过FeatureKey 与Feature关联），title 和 description， data（配置的数据，数据格式依据Feature配置模版生成，由用户填写），saData（键值对埋点数据，由用户自定义输入），圈人条件（包括开始时间、截止时间、圈人白名单、地理位置、灰度比例等）

新增后FeatureConfig记录状态默认为DRFAT。 管理课对DRAFT状态的数据执行发布操作，发布后记录变为PUBLISHED状态， 仅PUBLISHED状态的数据且符合圈人条件的配置C端可见。PUBLISHED状态的数据可执行撤销操作，撤销后记录变为DISABLED状态。

对DRAFT状态的记录执行删除操作后，记录会从数据库中物理删除，非DRAFT状态的记录，执行删除操作，进行逻辑删除状态标记为DELETED，DELETED状态的数据在管理界面上不可见。'''
    """

    response = agent_executor.invoke({"messages": [HumanMessage(content=message)]}, {'configurable': {"thread_id": "1"}})
    print(response)