"""
Graphiti 知识图谱客户端
替代 Zep Cloud，使用自托管的 Graphiti + Neo4j

Graphiti 是 async 库，本模块提供同步包装供 Flask 使用。
"""

import asyncio
import uuid
import time
from typing import Dict, Any, List, Optional
from datetime import datetime, timezone
from dataclasses import dataclass

from neo4j import GraphDatabase

from ..config import Config
from ..utils.logger import get_logger

logger = get_logger('foresight.graphiti_client')


def _safe_str(val):
    """Convert any value to JSON-safe string, handling Neo4j DateTime etc."""
    if val is None:
        return None
    return str(val)


from graphiti_core.embedder.client import EmbedderClient


class MiniMaxEmbedder(EmbedderClient):
    """
    Custom embedder for MiniMax API.
    Implements Graphiti's EmbedderClient interface.
    MiniMax uses 'texts' field instead of OpenAI's 'input' field.
    """

    def __init__(self, api_key: str, base_url: str = "https://api.minimax.chat/v1"):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")

    async def _call_api(self, texts: list[str]) -> list[list[float]]:
        import httpx
        async with httpx.AsyncClient(timeout=30) as client:
            response = await client.post(
                f"{self.base_url}/embeddings",
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
                json={"model": "embo-01", "texts": texts, "type": "db"},
            )
            response.raise_for_status()
            return response.json().get("vectors", [])

    async def create(self, input_data) -> list[float]:
        """Create embedding for a single text. Returns one vector."""
        if isinstance(input_data, str):
            texts = [input_data]
        elif isinstance(input_data, list) and input_data and isinstance(input_data[0], str):
            texts = [input_data[0]]
        else:
            texts = [str(input_data)]
        vectors = await self._call_api(texts)
        return vectors[0] if vectors else []

    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
        """Create embeddings for multiple texts."""
        return await self._call_api(input_data_list)


# Global event loop for async bridge
_loop: Optional[asyncio.AbstractEventLoop] = None


def _get_loop() -> asyncio.AbstractEventLoop:
    """Get or create a dedicated event loop for Graphiti async calls."""
    global _loop
    if _loop is None or _loop.is_closed():
        _loop = asyncio.new_event_loop()
    return _loop


def _run_async(coro):
    """Run an async coroutine synchronously."""
    loop = _get_loop()
    return loop.run_until_complete(coro)


@dataclass
class GraphitiNode:
    """Node data from Neo4j, compatible with Zep node format."""
    uuid_: str
    name: str
    labels: List[str]
    summary: str
    attributes: Dict[str, Any]
    created_at: Optional[str] = None

    @property
    def uuid(self):
        return self.uuid_


@dataclass
class GraphitiEdge:
    """Edge data from Neo4j, compatible with Zep edge format."""
    uuid_: str
    name: str
    fact: str
    source_node_uuid: str
    target_node_uuid: str
    attributes: Dict[str, Any]
    created_at: Optional[str] = None
    valid_at: Optional[str] = None
    invalid_at: Optional[str] = None
    expired_at: Optional[str] = None
    episodes: Optional[List[str]] = None

    @property
    def uuid(self):
        return self.uuid_


class GraphitiClient:
    """
    Graphiti + Neo4j 知识图谱客户端

    提供与原 Zep Cloud 兼容的接口:
    - create_graph / delete_graph
    - add_episodes (文本导入)
    - search (语义搜索)
    - get_all_nodes / get_all_edges
    - get_node / get_node_edges
    """

    _instance: Optional['GraphitiClient'] = None
    _graphiti = None
    _initialized = False

    def __init__(
        self,
        neo4j_uri: Optional[str] = None,
        neo4j_user: Optional[str] = None,
        neo4j_password: Optional[str] = None,
    ):
        self.neo4j_uri = neo4j_uri or Config.NEO4J_URI
        self.neo4j_user = neo4j_user or Config.NEO4J_USER
        self.neo4j_password = neo4j_password or Config.NEO4J_PASSWORD

        # Neo4j driver for direct queries
        self._driver = GraphDatabase.driver(
            self.neo4j_uri,
            auth=(self.neo4j_user, self.neo4j_password),
        )
        logger.info(f"GraphitiClient initialized: {self.neo4j_uri}")

    def _ensure_graphiti(self):
        """Lazy-init Graphiti (imports are heavy)."""
        if self._graphiti is not None:
            return

        from graphiti_core import Graphiti
        from graphiti_core.llm_client.openai_generic_client import OpenAIGenericClient
        from graphiti_core.llm_client.config import LLMConfig

        # Use SiliconFlow LLM for Graphiti (better structured output support)
        graphiti_api_key = Config.GRAPHITI_LLM_API_KEY or Config.LLM_API_KEY
        graphiti_base_url = Config.GRAPHITI_LLM_BASE_URL or Config.LLM_BASE_URL
        graphiti_model = Config.GRAPHITI_LLM_MODEL or Config.LLM_MODEL_NAME

        llm_config = LLMConfig(
            api_key=graphiti_api_key,
            model=graphiti_model,
            small_model=graphiti_model,
            base_url=graphiti_base_url,
        )
        llm_client = OpenAIGenericClient(config=llm_config)

        # Embedder: SiliconFlow free BAAI/bge-m3 (OpenAI-compatible)
        from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
        embedder_config = OpenAIEmbedderConfig(
            api_key=Config.EMBEDDING_API_KEY or Config.LLM_API_KEY,
            base_url=Config.EMBEDDING_BASE_URL or "https://api.siliconflow.cn/v1",
            embedding_model=Config.EMBEDDING_MODEL or "BAAI/bge-m3",
            embedding_dim=1024,
        )
        embedder = OpenAIEmbedder(config=embedder_config)

        # Reranker: use the LLM config (MiniMax)
        from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
        reranker = OpenAIRerankerClient(config=llm_config)

        self._graphiti = Graphiti(
            self.neo4j_uri,
            self.neo4j_user,
            self.neo4j_password,
            llm_client=llm_client,
            embedder=embedder,
            cross_encoder=reranker,
        )
        # Build indices
        _run_async(self._graphiti.build_indices_and_constraints())
        self._initialized = True
        logger.info("Graphiti core initialized with indices")

    @classmethod
    def get_instance(cls) -> 'GraphitiClient':
        """Singleton accessor."""
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    # ========== Graph CRUD ==========

    def create_graph(self, name: str) -> str:
        """Create a logical graph (just returns a group_id, Neo4j doesn't need explicit creation)."""
        graph_id = f"foresight_{uuid.uuid4().hex[:16]}"
        logger.info(f"Created graph group: {graph_id} ({name})")
        return graph_id

    def delete_graph(self, graph_id: str):
        """Delete all nodes and edges belonging to a graph group."""
        with self._driver.session() as session:
            # Delete edges first, then nodes
            session.run(
                "MATCH (a)-[r]-(b) WHERE r.group_id = $gid DELETE r",
                gid=graph_id,
            )
            session.run(
                "MATCH (n) WHERE n.group_id = $gid DETACH DELETE n",
                gid=graph_id,
            )
        logger.info(f"Deleted graph: {graph_id}")

    # ========== Episode ingestion ==========

    def add_episode(
        self,
        graph_id: str,
        text: str,
        source_description: str = "document",
        entity_types: Optional[Dict] = None,
        edge_types: Optional[Dict] = None,
    ):
        """Add a single text episode to the graph."""
        self._ensure_graphiti()

        kwargs = {
            "name": f"episode_{uuid.uuid4().hex[:8]}",
            "episode_body": text,
            "source_description": source_description,
            "reference_time": datetime.now(timezone.utc),
            "group_id": graph_id,
        }
        if entity_types:
            kwargs["entity_types"] = entity_types
        if edge_types:
            kwargs["edge_types"] = edge_types

        from graphiti_core.nodes import EpisodeType
        kwargs["source"] = EpisodeType.text

        _run_async(self._graphiti.add_episode(**kwargs))

    def add_episodes_batch(
        self,
        graph_id: str,
        texts: List[str],
        source_description: str = "document",
        entity_types: Optional[Dict] = None,
        edge_types: Optional[Dict] = None,
        progress_callback=None,
    ):
        """Add multiple text episodes using bulk API for parallel processing."""
        self._ensure_graphiti()
        total = len(texts)
        import time as _time

        # Use bulk API for parallel processing (much faster)
        BULK_SIZE = 5  # Process 5 episodes in parallel per batch

        for batch_start in range(0, total, BULK_SIZE):
            batch_texts = texts[batch_start:batch_start + BULK_SIZE]
            batch_end = min(batch_start + BULK_SIZE, total)

            if progress_callback:
                progress_callback(
                    f"正在并行处理第 {batch_start + 1}-{batch_end}/{total} 个文本块...",
                    batch_start / total,
                )

            start = _time.time()
            try:
                from graphiti_core.utils.bulk_utils import RawEpisode
                from graphiti_core.nodes import EpisodeType

                raw_episodes = [
                    RawEpisode(
                        name=f"episode_{batch_start + i}",
                        content=text,
                        source_description=source_description,
                        source=EpisodeType.text,
                        reference_time=datetime.now(timezone.utc),
                    )
                    for i, text in enumerate(batch_texts)
                ]

                _run_async(self._graphiti.add_episode_bulk(
                    bulk_episodes=raw_episodes,
                    group_id=graph_id,
                    entity_types=entity_types,
                    edge_types=edge_types,
                ))
            except Exception as e:
                logger.warning(f"Bulk failed for batch {batch_start}-{batch_end}, falling back to sequential: {e}")
                # Fallback: process one by one
                for i, text in enumerate(batch_texts):
                    try:
                        self.add_episode(graph_id, text, source_description, entity_types, edge_types)
                    except Exception as e2:
                        logger.error(f"Episode {batch_start + i + 1}/{total} failed: {e2}")
                        continue

            elapsed = _time.time() - start
            logger.info(f"Batch {batch_start + 1}-{batch_end}/{total} processed in {elapsed:.1f}s ({len(batch_texts)} episodes)")

            if progress_callback:
                progress_callback(
                    f"第 {batch_start + 1}-{batch_end}/{total} 处理完成（{len(batch_texts)}块用时 {elapsed:.0f}s）",
                    batch_end / total,
                )

    # ========== Search ==========

    def search(
        self,
        query: str,
        graph_id: Optional[str] = None,
        limit: int = 10,
    ) -> List[Dict[str, Any]]:
        """Search the graph for relevant edges/facts."""
        self._ensure_graphiti()

        kwargs = {"query": query, "num_results": limit}
        if graph_id:
            kwargs["group_ids"] = [graph_id]

        results = _run_async(self._graphiti.search(**kwargs))

        facts = []
        for edge in results:
            facts.append({
                "uuid": str(getattr(edge, 'uuid', '')),
                "fact": getattr(edge, 'fact', ''),
                "name": getattr(edge, 'name', ''),
                "source_node_uuid": str(getattr(edge, 'source_node_uuid', '')),
                "target_node_uuid": str(getattr(edge, 'target_node_uuid', '')),
            })
        return facts

    # ========== Node/Edge queries (direct Neo4j) ==========

    def get_all_nodes(self, graph_id: str, limit: int = 2000) -> List[GraphitiNode]:
        """Get all entity nodes for a graph group."""
        with self._driver.session() as session:
            result = session.run(
                """
                MATCH (n:Entity)
                WHERE n.group_id = $gid
                RETURN n, labels(n) as labels
                ORDER BY n.name
                LIMIT $limit
                """,
                gid=graph_id,
                limit=limit,
            )
            nodes = []
            for record in result:
                n = record["n"]
                raw_labels = record["labels"]
                # Filter internal labels
                labels = [l for l in raw_labels if l not in ("__Entity__",)]
                nodes.append(GraphitiNode(
                    uuid_=str(n.get("uuid", n.element_id)),
                    name=n.get("name", ""),
                    labels=labels,
                    summary=n.get("summary", ""),
                    attributes={k: _safe_str(v) for k, v in dict(n).items()} if n else {},
                    created_at=_safe_str(n.get("created_at")),
                ))
            return nodes

    def get_all_edges(self, graph_id: str) -> List[GraphitiEdge]:
        """Get all edges for a graph group."""
        with self._driver.session() as session:
            result = session.run(
                """
                MATCH (a)-[r]->(b)
                WHERE r.group_id = $gid
                RETURN r, type(r) as rtype, a.uuid as source_uuid, b.uuid as target_uuid
                """,
                gid=graph_id,
            )
            edges = []
            for record in result:
                r = record["r"]
                edges.append(GraphitiEdge(
                    uuid_=str(r.get("uuid", r.element_id)),
                    name=r.get("name", ""),
                    fact=r.get("fact", ""),
                    source_node_uuid=str(record["source_uuid"] or ""),
                    target_node_uuid=str(record["target_uuid"] or ""),
                    attributes={k: _safe_str(v) for k, v in dict(r).items()} if r else {},
                    created_at=_safe_str(r.get("created_at")),
                    valid_at=_safe_str(r.get("valid_at")),
                    invalid_at=_safe_str(r.get("invalid_at")),
                    expired_at=_safe_str(r.get("expired_at")),
                ))
            return edges

    def get_node(self, node_uuid: str) -> Optional[GraphitiNode]:
        """Get a single node by UUID."""
        with self._driver.session() as session:
            result = session.run(
                """
                MATCH (n:Entity {uuid: $uuid})
                RETURN n, labels(n) as labels
                """,
                uuid=node_uuid,
            )
            record = result.single()
            if not record:
                return None
            n = record["n"]
            labels = [l for l in record["labels"] if l not in ("__Entity__",)]
            return GraphitiNode(
                uuid_=str(n.get("uuid", n.element_id)),
                name=n.get("name", ""),
                labels=labels,
                summary=n.get("summary", ""),
                attributes={k: _safe_str(v) for k, v in dict(n).items()} if n else {},
                created_at=_safe_str(n.get("created_at")),
            )

    def get_node_edges(self, node_uuid: str) -> List[GraphitiEdge]:
        """Get all edges connected to a specific node."""
        with self._driver.session() as session:
            result = session.run(
                """
                MATCH (a)-[r]-(b)
                WHERE a.uuid = $uuid
                RETURN r,
                       CASE WHEN startNode(r) = a THEN a.uuid ELSE b.uuid END as source_uuid,
                       CASE WHEN startNode(r) = a THEN b.uuid ELSE a.uuid END as target_uuid
                """,
                uuid=node_uuid,
            )
            edges = []
            for record in result:
                r = record["r"]
                edges.append(GraphitiEdge(
                    uuid_=str(r.get("uuid", r.element_id)),
                    name=r.get("name", ""),
                    fact=r.get("fact", ""),
                    source_node_uuid=str(record["source_uuid"] or ""),
                    target_node_uuid=str(record["target_uuid"] or ""),
                    attributes=dict(r) if r else {},
                ))
            return edges

    def close(self):
        """Close connections."""
        if self._driver:
            self._driver.close()
        if self._graphiti:
            try:
                _run_async(self._graphiti.close())
            except Exception:
                pass
        logger.info("GraphitiClient closed")