"""Neo Cortex MCP — dual mode memory tools.

Modes (controlled by ENV vars):
    - Proxy (default): forwards to neo-cortex HTTP API on port 5074
    - Embedded: CORTEX_DATA_DIR set → local CortexService + LanceDB
    - Embedded + subscriber: CORTEX_DATA_DIR + MEMORY_HUB_URL set
      → also runs SignalR subscriber in daemon thread

Usage:
    python -m neo_cortex.mcp          # stdio transport (for Claude Code)
    python -m neo_cortex.mcp --http   # HTTP transport (for testing)
"""

# ============================================================
# BOOT LOGGING — raw file I/O before ANY third-party imports
# ============================================================
import os
import sys
import time

__version__ = "4.0.3"

_data_dir = os.environ.get("CORTEX_DATA_DIR")
_hub_url = os.environ.get("MEMORY_HUB_URL")
_embedded = _data_dir is not None

_log_dir = _data_dir or os.getcwd()
_log_path = os.path.join(_log_dir, "neo-cortex.log")
os.makedirs(_log_dir, exist_ok=True)


def _bl(msg: str) -> None:
    """Boot log — write directly to file before logging module is configured."""
    line = f"{time.strftime('%Y-%m-%d %H:%M:%S')} [BOOT] {msg}\n"
    try:
        with open(_log_path, "a", encoding="utf-8") as f:
            f.write(line)
    except Exception:
        pass


_bl(f"{'='*60}")
_bl(f"neo-cortex-mcp v{__version__} STARTING")
_bl(f"python: {sys.version}")
_bl(f"executable: {sys.executable}")
_bl(f"platform: {sys.platform}")
_bl(f"pid: {os.getpid()}")
_bl(f"cwd: {os.getcwd()}")
_bl(f"CORTEX_DATA_DIR: {_data_dir}")
_bl(f"MEMORY_HUB_URL: {_hub_url}")
_bl(f"JINA_API_KEY: {'SET' if os.environ.get('JINA_API_KEY') else 'NOT SET'}")
_bl(f"GROQ_API_KEY: {'SET' if os.environ.get('GROQ_API_KEY') else 'NOT SET'}")
_bl(f"embedded mode: {_embedded}")

# ============================================================
# IMPORTS — logged one by one to find crash point
# ============================================================
_bl("import: logging")
import logging

_bl("import: threading")
import threading

_bl("import: pathlib.Path")
from pathlib import Path

_bl("import: typing.Optional")
from typing import Optional

_bl("import: httpx")
try:
    import httpx
    _bl("import: httpx OK")
except Exception as e:
    _bl(f"import: httpx FAILED: {e}")
    raise

_bl("import: fastmcp.FastMCP")
try:
    from fastmcp import FastMCP
    _bl("import: fastmcp OK")
except Exception as e:
    _bl(f"import: fastmcp FAILED: {e}")
    raise

# ============================================================
# LOGGER setup (now that logging is imported)
# ============================================================
logger = logging.getLogger("neo-cortex-mcp")

# ============================================================
# MODE INIT — embedded or proxy
# ============================================================
_cortex = None
_http = None
_init_error = None

if _embedded:
    _bl("embedded: importing subscriber.configure_paths")
    from neo_cortex.subscriber import configure_paths
    _bl("embedded: configure_paths()")
    configure_paths(_data_dir)
    Path(_data_dir).mkdir(parents=True, exist_ok=True)

    try:
        _bl("embedded: importing neo_cortex.config")
        from neo_cortex import config
        _bl(f"embedded: CORTEX_DB_PATH={config.CORTEX_DB_PATH}")
        _bl(f"embedded: MEMORY_INDEX_DB_PATH={config.MEMORY_INDEX_DB_PATH}")

        _bl("embedded: importing GroqClassifier")
        from neo_cortex.classifier import GroqClassifier

        _bl("embedded: importing CortexService")
        from neo_cortex.cortex import CortexService

        _bl("embedded: importing JinaEmbedder")
        from neo_cortex.embedder import JinaEmbedder

        _bl("embedded: importing MemoryIndex")
        from neo_cortex.memory_index import MemoryIndex

        _bl("embedded: importing MemoryStore")
        from neo_cortex.store import MemoryStore

        _bl(f"embedded: creating MemoryStore(db_path={config.CORTEX_DB_PATH})")
        store = MemoryStore(config.CORTEX_DB_PATH, config.COLLECTION_NAME)
        _bl(f"embedded: MemoryStore OK, count={store.count()}")

        _bl("embedded: load_api_key('jina')")
        jina_key = config.load_api_key("jina")
        _bl(f"embedded: jina key loaded ({len(jina_key)} chars)")

        _bl(f"embedded: platform={sys.platform}, verify_ssl={sys.platform != 'win32'}")

        _bl("embedded: creating JinaEmbedder")
        embedder = JinaEmbedder(jina_key)
        _bl("embedded: JinaEmbedder OK")

        _bl("embedded: load_api_key('groq')")
        groq_key = config.load_api_key("groq")
        _bl(f"embedded: groq key loaded ({len(groq_key)} chars)")

        _bl("embedded: creating GroqClassifier")
        classifier = GroqClassifier(groq_key)
        _bl("embedded: GroqClassifier OK")

        _bl("embedded: creating MemoryIndex")
        index = MemoryIndex(config.MEMORY_INDEX_DB_PATH)
        _bl(f"embedded: MemoryIndex OK, count={index.count()}")

        _bl("embedded: creating CortexService")
        _cortex = CortexService(store, embedder, classifier, index=index)
        _bl("embedded: CortexService READY")

    except Exception as e:
        _init_error = str(e)
        _bl(f"embedded: INIT FAILED: {e}")
        import traceback
        _bl(traceback.format_exc())
        _cortex = None
else:
    _bl("proxy: importing config")
    from neo_cortex import config
    CORTEX_BASE = f"http://127.0.0.1:{config.CORTEX_PORT}"
    _bl(f"proxy: connecting to {CORTEX_BASE}")
    _http = httpx.Client(base_url=CORTEX_BASE, timeout=15.0)
    _bl("proxy: httpx client OK")


def _check_embedded():
    """Raise clear error if embedded mode but cortex init failed."""
    if _embedded and _cortex is None:
        raise RuntimeError(f"CortexService init failed: {_init_error}. Check JINA_API_KEY and GROQ_API_KEY env vars.")


def _get(path: str, **params) -> dict:
    _check_embedded()
    resp = _http.get(path, params={k: v for k, v in params.items() if v is not None})
    resp.raise_for_status()
    return resp.json()


def _post(path: str, data: dict = None) -> dict:
    _check_embedded()
    resp = _http.post(path, json=data)
    resp.raise_for_status()
    return resp.json()


# --- Subscriber thread (embedded + hub_url) ---
def _start_subscriber(hub_url: str, cortex) -> threading.Thread:
    """Start SignalR subscriber in a daemon thread with its own asyncio loop."""
    import asyncio as _asyncio

    from neo_cortex.subscriber import TurnAccumulator

    async def _run():
        from pysignalr.client import SignalRClient

        client = SignalRClient(hub_url)
        accumulator = TurnAccumulator()

        async def on_user_message(args):
            data = args[0] if args else {}
            accumulator.on_user_message(data.get("sessionId", ""), data.get("text", ""))

        async def on_assistant_message(args):
            data = args[0] if args else {}
            accumulator.on_assistant_message(
                data.get("sessionId", ""), data.get("text", ""), data.get("model", "unknown"))

        async def on_tool_use(args):
            data = args[0] if args else {}
            accumulator.on_tool_use(data.get("sessionId", ""), data.get("toolName", ""))

        async def on_tool_result(args):
            logger.debug("ToolResult received")

        async def on_turn_complete(args):
            if not accumulator.is_complete:
                accumulator.reset()
                return
            from neo_cortex.models import IngestRequest
            req = 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(req)
                if memory_id:
                    logger.info("Subscriber ingested: %s (turn %d)", memory_id, accumulator.turn_number)
            except Exception:
                logger.exception("Subscriber ingest failed for turn %d", accumulator.turn_number)
            accumulator.reset()

        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)

        logger.info("Subscriber connecting to %s", hub_url)
        await client.run()

    def _thread_target():
        _asyncio.run(_run())

    t = threading.Thread(target=_thread_target, daemon=True, name="cortex-subscriber")
    t.start()
    return t


# --- MCP server ---
_bl("creating FastMCP server")
mcp = FastMCP(
    "neo-cortex",
    instructions="Neo's memory cortex. Semantic memory search, ingestion, dream cycles, and timeline.",
)
_bl("FastMCP server created")


@mcp.tool()
async def memory_query(query: str, n: int = 5) -> dict:
    """Search past conversation memories by semantic similarity.

    Args:
        query: What to search for (natural language)
        n: Number of results to return (default 5)
    """
    if _cortex:
        result = await _cortex.recall(query, n=n, smart=True)
        return result.model_dump() if hasattr(result, "model_dump") else result.__dict__
    return _get("/api/query", q=query, n=n, smart="true")


@mcp.tool()
async def memory_stats() -> dict:
    """Get cortex statistics: total memories, sessions, energy levels, dream count."""
    if _cortex:
        result = _cortex.stats()
        return result.model_dump() if hasattr(result, "model_dump") else result.__dict__
    return _get("/api/stats")


@mcp.tool()
async def memory_dream() -> dict:
    """Run one dream cycle: boost high-energy memories, decay others."""
    if _cortex:
        result = await _cortex.dream()
        return result.model_dump() if hasattr(result, "model_dump") else result.__dict__
    return _post("/api/dream")


@mcp.tool()
async def memory_ingest(turn: dict) -> dict:
    """Manually ingest a conversation turn into memory.

    Args:
        turn: JSON with session_id, question, answer, model, stats
    """
    if _cortex:
        from neo_cortex.models import IngestRequest
        req = IngestRequest(**turn)
        memory_id = await _cortex.ingest(req)
        return {"status": "ok", "id": memory_id}
    return _post("/api/ingest", data=turn)


@mcp.tool()
async def memory_timeline(n: int = 10, project: Optional[str] = None) -> dict:
    """Get most recent memories by timestamp (newest first).

    Use this at session start to load recent context, or to see what happened lately.

    Args:
        n: Number of recent memories to return (default 10)
        project: Optional filter by project name
    """
    if _cortex:
        result = await _cortex.timeline(n=n, project=project)
        return result.model_dump() if hasattr(result, "model_dump") else result.__dict__
    return _get("/api/timeline", n=n, project=project)


@mcp.tool()
async def memory_search(query: str = "", project: Optional[str] = None,
                        activity: Optional[str] = None, n: int = 10) -> dict:
    """Search memory index. Returns compact results (title/project/date).
    Use memory_get for full details of specific memories.

    Args:
        query: Search text (or empty for recent)
        project: Filter by project name
        activity: Filter by activity type
        n: Max results
    """
    if _cortex:
        if hasattr(_cortex, '_index') and _cortex._index:
            results = _cortex._index.search_fts(query=query, project=project, activity=activity, n=n)
            return {"results": [r.__dict__ for r in results]}
        # Fallback: use recall
        result = await _cortex.recall(query or "recent", n=n)
        return result.model_dump() if hasattr(result, "model_dump") else result.__dict__
    return _get("/api/search", q=query, project=project, activity=activity, n=n)


@mcp.tool()
async def memory_get(ids: str) -> dict:
    """Get full memory details by IDs (comma-separated).

    Args:
        ids: Comma-separated memory IDs (e.g. "v3_abc_1_12345,v3_def_2_67890")
    """
    if _cortex:
        if hasattr(_cortex, '_index') and _cortex._index:
            id_list = [i.strip() for i in ids.split(",")]
            records = _cortex._index.get_by_ids(id_list)
            return {"memories": [r.__dict__ for r in records]}
        return {"error": "memory_get requires SQLite index in embedded mode"}
    return _get("/api/memories", ids=ids)


# --- Entrypoint ---


def main():
    _bl("main() called")

    # Configure structured logging (FileHandler + StreamHandler)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
        handlers=[
            logging.FileHandler(_log_path, encoding="utf-8"),
            logging.StreamHandler(),
        ],
    )

    # Diagnostic summary
    import platform
    try:
        import lancedb as _lancedb
        lance_ver = _lancedb.__version__
    except Exception:
        lance_ver = "N/A"

    logger.info(
        "neo-cortex-mcp v%s | python %s | lancedb %s | os=%s | embedded=%s | data_dir=%s",
        __version__, platform.python_version(), lance_ver,
        platform.system(), _embedded, _data_dir,
    )
    if _cortex:
        logger.info("CortexService OK")
    else:
        logger.error("CortexService FAILED: %s", _init_error)

    # Start subscriber if both env vars set
    if _hub_url and _cortex:
        _start_subscriber(_hub_url, _cortex)

    _bl("starting MCP transport")
    if "--http" in sys.argv:
        mcp.run(transport="streamable-http", host="127.0.0.1", port=5075)
    else:
        mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
