"""SignalR subscriber for neo-reloaded Memory Hub.

Connects to /hub/memory, accumulates events per turn,
and ingests completed turns into the cortex storage.

Usage:
    python -m neo_cortex.subscriber --hub-url http://localhost:5070/hub/memory --data-dir /path/to/data
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger("neo_cortex.subscriber")


@dataclass
class TurnAccumulator:
    """Accumulates SignalR memory events within a single turn."""

    session_id: str = ""
    question: str = ""
    answer: str = ""
    model: str = "unknown"
    tools_used: list[str] = field(default_factory=list)
    turn_number: int = 0

    def on_user_message(self, session_id: str, text: str) -> None:
        self.session_id = session_id
        self.question = text
        self.turn_number += 1

    def on_assistant_message(self, session_id: str, text: str, model: str) -> None:
        self.session_id = session_id
        self.answer = text
        self.model = model

    def on_tool_use(self, session_id: str, tool_name: str) -> None:
        self.session_id = session_id
        if tool_name not in self.tools_used:
            self.tools_used.append(tool_name)

    @property
    def is_complete(self) -> bool:
        return bool(self.question and self.answer)

    def reset(self) -> None:
        self.question = ""
        self.answer = ""
        self.tools_used = []


def configure_paths(data_dir: str) -> None:
    """Override cortex env vars so storage goes to data_dir/cortex_db/."""
    db_dir = Path(data_dir) / "cortex_db"
    os.environ["CORTEX_DB_PATH"] = str(db_dir)
    os.environ["MEMORY_INDEX_DB_PATH"] = str(db_dir / "memory_index.db")
    os.environ["CONVERSATION_DB_PATH"] = str(db_dir / "conversation_log.db")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Neo Cortex SignalR subscriber")
    parser.add_argument("--hub-url", required=True, help="SignalR hub URL (e.g. http://localhost:5070/hub/memory)")
    parser.add_argument("--data-dir", required=True, help="Directory for cortex data (ChromaDB + SQLite)")
    return parser.parse_args(argv)


async def run_subscriber(hub_url: str, data_dir: str) -> None:
    """Connect to SignalR hub and ingest memory events."""
    from pysignalr.client import SignalRClient

    # Configure paths before importing cortex (uses env vars at import time)
    configure_paths(data_dir)
    Path(data_dir).mkdir(parents=True, exist_ok=True)

    # Lazy import after path config
    from neo_cortex import config
    from neo_cortex.classifier import GroqClassifier
    from neo_cortex.cortex import CortexService
    from neo_cortex.embedder import JinaEmbedder
    from neo_cortex.memory_index import MemoryIndex
    from neo_cortex.store import MemoryStore

    jina_key = config.load_api_key("jina")
    groq_key = config.load_api_key("groq")

    store = MemoryStore(config.CORTEX_DB_PATH, config.COLLECTION_NAME)
    embedder = JinaEmbedder(jina_key)
    classifier = GroqClassifier(groq_key)
    index = MemoryIndex(config.MEMORY_INDEX_DB_PATH)

    cortex = CortexService(store, embedder, classifier, index=index)
    accumulator = TurnAccumulator()

    logger.info("CortexService initialized: db=%s, index=%s",
                config.CORTEX_DB_PATH, config.MEMORY_INDEX_DB_PATH)

    client = SignalRClient(hub_url)

    async def on_open() -> None:
        logger.info("Connected to %s", hub_url)

    async def on_close() -> None:
        logger.warning("Disconnected from %s", hub_url)

    async def on_error(message: Any) -> None:
        logger.error("SignalR error: %s", message)

    async def on_user_message(args: list[Any]) -> None:
        data = args[0] if args else {}
        accumulator.on_user_message(
            data.get("sessionId", ""),
            data.get("text", ""),
        )
        logger.debug("UserMessage: session=%s", accumulator.session_id)

    async def on_assistant_message(args: list[Any]) -> None:
        data = args[0] if args else {}
        accumulator.on_assistant_message(
            data.get("sessionId", ""),
            data.get("text", ""),
            data.get("model", "unknown"),
        )
        logger.debug("AssistantMessage: session=%s, model=%s", accumulator.session_id, accumulator.model)

    async def on_tool_use(args: list[Any]) -> None:
        data = args[0] if args else {}
        accumulator.on_tool_use(
            data.get("sessionId", ""),
            data.get("toolName", ""),
        )
        logger.debug("ToolUse: %s", data.get("toolName", ""))

    async def on_tool_result(args: list[Any]) -> None:
        # Tool results are tracked but not accumulated separately
        logger.debug("ToolResult received")

    async def on_turn_complete(args: list[Any]) -> None:
        if not accumulator.is_complete:
            logger.debug("TurnComplete but no complete Q&A, skipping ingest")
            accumulator.reset()
            return

        from neo_cortex.models import IngestRequest

        request = IngestRequest(
            session_id=accumulator.session_id,
            question=accumulator.question,
            answer=accumulator.answer,
            model=accumulator.model,
            tools_used=accumulator.tools_used,
            turn_number=accumulator.turn_number,
        )

        try:
            memory_id = await cortex.ingest(request)
            if memory_id:
                logger.info("Ingested memory: %s (turn %d)", memory_id, accumulator.turn_number)
            else:
                logger.debug("Filtered as noise (turn %d)", accumulator.turn_number)
        except Exception:
            logger.exception("Ingest failed for turn %d", accumulator.turn_number)

        accumulator.reset()

    async def on_error_event(args: list[Any]) -> None:
        data = args[0] if args else {}
        logger.warning("MemoryError: %s", data.get("message", "unknown"))

    # Register handlers
    client.on_open(on_open)
    client.on_close(on_close)
    client.on_error(on_error)
    client.on("MemoryUserMessage", on_user_message)
    client.on("MemoryAssistantMessage", on_assistant_message)
    client.on("MemoryToolUse", on_tool_use)
    client.on("MemoryToolResult", on_tool_result)
    client.on("MemoryTurnComplete", on_turn_complete)
    client.on("MemoryError", on_error_event)

    logger.info("Starting subscriber: hub=%s, data=%s", hub_url, data_dir)
    await client.run()


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    args = parse_args()
    asyncio.run(run_subscriber(args.hub_url, args.data_dir))


if __name__ == "__main__":
    main()
