from langchain_core.retrievers import BaseRetriever
from langchain.schema.document import Document
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun

from qwin_chat_bot_domain import extract_parse_json
from qdrant import QwinQdrantClient
from db import QwinCasinoDao
from finance_info_provider import DepositInfoProvider, WithdrawalInfoProvider

class FinanceRetriever(BaseRetriever):
    user_id: str
    qdrant: QwinQdrantClient
    db: QwinCasinoDao
    deposit_info_provider: DepositInfoProvider
    withdrawal_info_provider: WithdrawalInfoProvider
    
    def _get_relevant_documents(self, query_str, run_manager: CallbackManagerForRetrieverRun):
        query = extract_parse_json(query_str)
        category = query['category']
        
        if category == 'terms':
            return [Document(page_content=doc) 
                    for doc in self.qdrant.search_finance_docs()]
        
        if category == 'deposit':
            return self.deposit_info_provider.deposit_info(self.user_id, query)
        
        if category == 'withdrawal':
            return self.withdrawal_info_provider.withdrawal_info(self.user_id, query)
        
        return [Document(page_content='Ask the user to clarify what he wants to know')]
        
        
            
def create_finance_retriever(user_id: str) -> FinanceRetriever:
    deposit_info_provider = DepositInfoProvider(qdrant=QwinQdrantClient(), db=QwinCasinoDao())
    withdrawal_info_provider = WithdrawalInfoProvider(qdrant=QwinQdrantClient(), db=QwinCasinoDao())
    
    return FinanceRetriever(user_id=user_id, 
                            qdrant=QwinQdrantClient(), 
                            db=QwinCasinoDao(), 
                            deposit_info_provider=deposit_info_provider,
                            withdrawal_info_provider=withdrawal_info_provider)