import json
from typing import Dict, Any, Tuple, List, Optional
from alayaflow.clients.alayamem.http_client import HttpAlayaMemClient
from alayaflow.clients.alayamem.types import HybridQueryResponse, HybridQueryResults
from alayaflow.utils.logger import AlayaFlowLogger


logger = AlayaFlowLogger()

class RetrieveComponent:
    def __init__(self, client: HttpAlayaMemClient):
        self.client = client
    
    def __call__(self, query: str, limit: int = 5, user_id: str = None) -> List[str]:
        try:
            if not user_id:
                logger.info(f"RetrieveComponent: user_id is None, query: {query}")
                return []
            result = self.client.query_file(user_id=user_id, message=query, limit=limit)
            # query_file 返回格式: {'msg': 'query_file', 'results': {'documents': [['doc1', 'doc2', ...]], ...}, ...}
            if isinstance(result, dict):
                results = result.get('results', {})
                if isinstance(results, dict):
                    documents = results.get('documents', [])
                    # documents 格式是 [['doc1', 'doc2', ...]]，需要取第一个元素
                    if documents and isinstance(documents, list) and len(documents) > 0:
                        if isinstance(documents[0], list):
                            return documents[0]
                        return documents
            return []
        except Exception as e:
            logger.error(f"RetrieveComponent: exception: {type(e).__name__}: {e}")
            return []


class HybridRetrieveComponent:
    """混合检索组件，结合向量检索和关键词检索，返回检索结果"""
    def __init__(self, client: HttpAlayaMemClient):
        self.client = client
    
    def __call__(self, query: str, limit: int = 5, user_id: str = None) -> Optional[HybridQueryResults]:
        """
        执行混合检索，返回检索结果
        
        Args:
            query: 查询文本
            limit: 返回结果数量，默认 5
            user_id: 用户ID
            
        Returns:
            HybridQueryResults: 包含以下字段：
                - ids: 文档ID列表 (List[List[str]])
                - documents: 文档内容列表 (List[List[str]])
                - metadatas: 元数据列表 (List[List[Dict[str, Any]]])
                - distances: 距离列表 (List[List[float]])
        """
        try:
            if not user_id:
                logger.info(f"HybridRetrieveComponent: user_id is None, query: {query}")
                return None
            response: HybridQueryResponse = self.client.hybrid_query(query=query, user_id=user_id, limit=limit)
            return response.get('results')
        except Exception as e:
            logger.error(f"HybridRetrieveComponent: exception: {type(e).__name__}: {e}")
            return None


class MemoryRetrieveComponent:
    def __init__(self, client: HttpAlayaMemClient):
        self.client = client

    def __call__(self, user_id: str, profile_id: str) -> str:
        try:
            memory_obj = self.client.get_profile(user_id=user_id, profile_id=profile_id)
            # 提取 content 部分：返回格式 {'code': 0, 'msg': 'success', 'data': {'content': {...}, ...}}
            if isinstance(memory_obj, dict):
                data = memory_obj.get('data', {})
                if isinstance(data, dict):
                    content = data.get('content', {})
                    return json.dumps(content, ensure_ascii=False)
            return json.dumps(memory_obj, ensure_ascii=False)
        except Exception:
            return ""


class MemoryUpdateComponent:
    def __init__(self, client: HttpAlayaMemClient):
        self.client = client

    def __call__(self, user_id: str, patch: Dict[str, Any], profile_id: str) -> Tuple[str, bool]:
        try:
            resp = self.client.update_profile(user_id=user_id, profile_id=profile_id, content=patch)
            return json.dumps(resp or {}, ensure_ascii=False), True
        except Exception:
            return "", False
