from langchain import LLMChain, PromptTemplate
from langchain.agents import tool

from models import llm
from agents.tools.smart_domain.common import getPrefix
from agents.tools.smart_domain.db_entity_repository import db_entity_architecture, db_entity_test_strategy
from agents.tools.smart_domain.association_impl import association_impl_architecture, association_impl_test_strategy


persistent_task = """"Your task is to generate the persistent layer tests and product code."""
persistent_tech_stack = """Java17、reactor、lombok、Junit5、reactor test、Mockito、 Spring Data Reactive Couchbase、Testcontainers、Couchbase、WebClient"""
persistent_architecture = f"""the persistent layer inclue 3 componets:
{db_entity_architecture}
{association_impl_architecture}"""

persistent_test_strategy = f"""{db_entity_test_strategy}
{association_impl_test_strategy}"""

PERSISTENT_LAYER = getPrefix(persistent_task, persistent_tech_stack, persistent_architecture, persistent_test_strategy) + """

Use the following format:
request: the request that you need to fulfill include Entity and Association of domain layer

DBEntity:
```
the DBEntity code that you write to fulfill the request, follow TechStack and Architecture
```

Repository:
```
the Repository code that you write to fulfill the request, follow TechStack and Architecture
```

Association Impletation:
```
the Association Impletation code that you write to fulfill the request, follow TechStack and Architecture
```

Test:
```
the test code that you write to fulfill the request, follow TechStack Architecture and TestStrategy
```

request: {input}"""

PERSISTENT_LAYER_PROMPT = PromptTemplate(input_variables=["input"], template=PERSISTENT_LAYER,)

persistentChain = LLMChain(llm = llm(temperature=0.1), prompt=PERSISTENT_LAYER_PROMPT)


@tool("Generate Persistent Layer Code", return_direct=True)
def persistentLayerCodeGenerator(input: str) -> str:
    '''useful for when you need to generate persistent layer code'''
    response = persistentChain.run(input)
    return response