#!/usr/bin/env python3
# pyright: reportCallIssue=false
"""
A FastAPI application for Memgraph with a web interface.

Note: FastMCP 3.x decorated functions are callable at runtime but Pylance
type inference may show them as FunctionTool. The pyright directive above
suppresses these false positives.
"""

from collections.abc import AsyncGenerator
import atexit
from contextlib import asynccontextmanager
from typing import Any
import logging
import os
import webbrowser

from starlette.types import Lifespan
from fastapi import FastAPI
from fastmcp import FastMCP

from memgraph.config import IN_DOCKER, Config, default_config_dir, load_config
from memgraph.sqlite_backend import SQLiteKnowledgeGraphDB

from memgraph.launcher import save_server_info

logger = logging.getLogger(__name__)

# Module-level mcp instance for testing and imports
# Initialized with default config, can be overridden
_default_mcp: FastMCP | None = None


def get_mcp(config: Config | None = None) -> FastMCP:
    """
    Get or create the MCP instance.

    Args:
        config: Optional configuration to use. If None, uses default config.

    Returns:
        FastMCP instance
    """
    global _default_mcp
    if _default_mcp is None or config is not None:
        _default_mcp = setup_mcp(config or load_config(default_config_dir()))
    return _default_mcp


def setup_mcp(config: Config) -> FastMCP:
    """
    Set up the FastMCP application with Memgraph knowledge graph tools.

    Returns:
        FastMCP: Configured FastMCP application
    """
    mcp = FastMCP(
        name="Zabob Memgraph Knowledge Graph Server",
        instructions="A FastAPI application for Memgraph with a web interface.",
        lifespan=get_lifespan_hook(config),
    )
    DB = SQLiteKnowledgeGraphDB(config)

    @mcp.tool
    async def read_graph(name: str = "default") -> dict[str, Any]:
        """
        Read the complete knowledge graph from the database.

        This returns all entities, relations, and observations in the graph,
        formatted for visualization or analysis.

        Args:
            name (str): Graph identifier (default: 'default')

        Returns:
            dict: Complete graph data with entities, relations, and observations
        """
        logger.info(f"Reading graph: {name}")
        return await DB.read_graph()

    @mcp.tool
    async def search_nodes(query: str) -> dict[str, Any]:
        """
        Search the knowledge graph for entities and relations matching the query.

        Performs full-text search across entity names, types, and observations.

        Args:
            query (str): Search query string

        Returns:
            dict: Search results containing matching entities and their metadata
        """
        logger.info(f"Searching graph with query: {query}")
        return await DB.search_nodes(query)

    @mcp.tool
    async def get_server_info() -> dict[str, Any]:
        """
        Get information about this server instance.

        Returns server identity information including name, version, port, host,
        database path, and container details if running in Docker.
        Useful for distinguishing between multiple server instances in multi-agent scenarios.

        Returns:
            dict: Server information with name, version, port, host, database_path,
                  in_docker, and container_name (if applicable)
        """
        from memgraph.__version__ import __version__

        info = {
            "name": config.get("name", "default"),
            "version": __version__,
            "port": config.get("real_port") if IN_DOCKER else config.get("port"),
            "host": config.get("host"),
            "database_path": str(config.get("database_path")),
            "in_docker": IN_DOCKER,
        }

        if IN_DOCKER:
            info["container_name"] = config.get("container_name")

        logger.info(f"Returning server info for '{info['name']}'")
        return info

    @mcp.tool
    async def get_stats() -> dict[str, Any]:
        """
        Get statistics about the knowledge graph.

        Returns counts and metadata about entities, relations, and observations
        in the database.

        Returns:
            dict: Statistics including entity count, relation count, observation count, etc.
        """
        logger.info("Getting graph statistics")
        return await DB.get_stats()

    @mcp.tool
    async def create_entities(entities: list[dict[str, Any]]) -> dict[str, Any]:
        """
        Create new entities in the knowledge graph.

        **IMPORTANT**: Always search first to avoid duplicates!
        - Use search_nodes() before creating to check if entity already exists
        - Creating duplicates causes relation dilution (relationships scattered across copies)
        - Search by name, type, or related concepts before creating new entities

        Each entity should have:
        - name (str): Entity identifier
        - entityType (str): Type of entity
        - observations (list[str], optional): Initial observations

        Args:
            entities (list[dict]): List of entity objects to create

        Returns:
            dict: Result with count of entities created
        """
        logger.info(f"Creating {len(entities)} entities")
        await DB.create_entities(entities)
        return {"created": len(entities), "entities": [e.get("name") for e in entities]}

    @mcp.tool
    async def create_relations(relations: list[dict[str, Any]], external_refs: list[str]) -> dict[str, Any]:
        """
        Create new relations between entities in the knowledge graph.

        Each relation should have:
        - source (str): Source entity name
        - target (str): Target entity name
        - relation (str): Type of relation

        Args:
            relations (list[dict]): List of relation objects to create
            external_refs (list[str]): Entity names that must exist (REQUIRED).
                Validates all referenced entities exist before creating relations.
                Returns error if any are missing. Use create_subgraph if you need
                to create entities and relations together.

        Returns:
            dict: Result with count of relations created, or error if validation fails
        """
        logger.info(f"Creating {len(relations)} relations (external_refs: {external_refs})")

        # Map field names from MCP format to SQLite backend format
        # Handle both formats: MCP (source/target/relation) and backend (from_entity/to/relationType)
        mapped_relations = []
        for r in relations:
            if "source" in r:
                # MCP format
                mapped_relations.append(
                    {
                        "from_entity": r["source"],
                        "to": r["target"],
                        "relationType": r["relation"],
                    }
                )
            else:
                # Already in backend format
                mapped_relations.append(r)

        # Get initial count for verification
        initial_stats = await DB.get_stats()
        initial_count = initial_stats.get("relation_count", 0)

        # Create relations with validation
        try:
            await DB.create_relations(mapped_relations, external_refs=external_refs)
        except ValueError as e:
            logger.error(f"Validation failed: {e}")
            return {"error": str(e), "created": 0, "relations": []}

        # Verify they were created
        final_stats = await DB.get_stats()
        final_count = final_stats.get("relation_count", 0)
        actual_created = final_count - initial_count

        if actual_created != len(relations):
            logger.warning(f"Expected to create {len(relations)} relations, but only {actual_created} were created")

        return {
            "created": actual_created,
            "relations": [f"{r.get('source')} -> {r.get('target')}" for r in relations],
        }

    @mcp.tool
    async def create_subgraph(
        entities: list[dict[str, Any]],
        relations: list[dict[str, Any]],
        external_refs: list[str] | None = None,
        observations: dict[str, list[str]] | None = None,
    ) -> dict[str, Any]:
        """
        Create a subgraph atomically with entities, relations, and observations.

        **IMPORTANT**: Always search first to avoid duplicates!
        - Use search_nodes() before creating to check if entities already exist
        - Creating duplicates causes relation dilution (relationships scattered across copies)
        - Search for similar concepts and add to existing entities instead of creating new ones
        - Use external_refs to link to existing entities rather than duplicating them

        This is a high-level operation that combines entity creation, observation
        addition, and relation creation in a single atomic transaction. Use this
        when you need to add a complete, self-contained graph pattern.

        Args:
            entities (list[dict]): New entities to create. Each should have:
                - name (str): Entity name
                - entityType (str): Entity type
                - observations (list[str], optional): Initial observations
            relations (list[dict]): Relations to create. Each should have:
                - source (str): Source entity name
                - target (str): Target entity name
                - relation (str): Relation type
            external_refs (list[str], optional): Existing entity names being referenced.
                These entities must already exist. Defaults to empty list.
            observations (dict[str, list[str]], optional): Additional observations to add
                to any entity (new or existing). Keys are entity names, values are
                lists of observation strings.

        Returns:
            dict: Result with counts of created entities, relations, and observations added,
                  or error if validation fails

        Example:
            create_subgraph(
                entities=[{"name": "Task-123", "entityType": "task", "observations": ["Started today"]}],
                external_refs=["Bob Kerns", "zabob-memgraph"],
                observations={
                    "Task-123": ["Assigned to Bob"],
                    "Bob Kerns": ["Working on Task-123"]
                },
                relations=[
                    {"source": "Task-123", "target": "zabob-memgraph", "relation": "modifies"},
                    {"source": "Bob Kerns", "target": "Task-123", "relation": "assigned_to"}
                ]
            )
        """
        logger.info(
            f"Creating subgraph: {len(entities)} entities, {len(relations)} relations, "
            f"{len(observations or {})} observation groups, external_refs: {external_refs}"
        )

        # Map relation formats
        mapped_relations = []
        for r in relations:
            if "source" in r:
                mapped_relations.append(
                    {
                        "from_entity": r["source"],
                        "to": r["target"],
                        "relationType": r["relation"],
                    }
                )
            else:
                mapped_relations.append(r)

        try:
            await DB.create_subgraph(
                entities=entities,
                relations=mapped_relations,
                external_refs=external_refs,
                observations=observations,
            )

            return {
                "created_entities": len(entities),
                "created_relations": len(relations),
                "observation_groups": len(observations or {}),
                "entities": [e["name"] for e in entities],
                "relations": [f"{r.get('source')} -> {r.get('target')}" for r in relations],
            }
        except ValueError as e:
            logger.error(f"Subgraph creation failed: {e}")
            return {"error": str(e), "created_entities": 0, "created_relations": 0, "observation_groups": 0}

    @mcp.tool
    async def add_observations(
        entity_name: str, observations: list[str], external_refs: list[str] | None = None
    ) -> dict[str, Any]:
        """
        Add observations to an existing entity.

        Args:
            entity_name (str): Name of the entity to add observations to
            observations (list[str]): List of observation strings to add
            external_refs (list[str], optional): Entity names that must exist.
                If provided, validates all referenced entities exist.
                Defaults to [entity_name] if not specified.

        Returns:
            dict: Result with count of observations added, or error if validation fails
        """
        logger.info(f"Adding {len(observations)} observations to {entity_name} (external_refs: {external_refs})")

        # Default to validating the target entity exists
        if external_refs is None:
            external_refs = [entity_name]

        # Validate entity exists
        try:
            import sqlite3

            with sqlite3.connect(DB.db_path) as conn:
                conn.row_factory = sqlite3.Row
                placeholders = ",".join("?" * len(external_refs))
                cursor = conn.execute(
                    f"SELECT name FROM entities WHERE name IN ({placeholders})",
                    external_refs,
                )
                found = {row["name"] for row in cursor}
                missing = set(external_refs) - found
                if missing:
                    error_msg = f"Referenced entities not found: {sorted(missing)}"
                    logger.error(error_msg)
                    return {"error": error_msg, "entity": entity_name, "added": 0}
        except Exception as e:
            logger.error(f"Validation failed: {e}")
            return {"error": str(e), "entity": entity_name, "added": 0}

        # Create a pseudo-entity update with new observations
        await DB.create_entities(
            [
                {
                    "name": entity_name,
                    "entityType": "update",  # Will merge with existing
                    "observations": observations,
                }
            ]
        )
        return {"entity": entity_name, "added": len(observations)}

    @mcp.tool
    async def open_browser(node_id: str | None = None) -> dict[str, Any]:
        """
        Open a browser window to visualize the knowledge graph.

        Reads the server URL from server_info.json or scans for running servers.
        Optionally focuses on a specific node if node_id is provided.

        If multiple servers are running, opens the first one found.

        Note: Only available when running locally, not in Docker containers.

        Args:
            node_id (str, optional): ID of a specific node to focus on in the visualization

        Returns:
            dict: Status of the operation with URL that was opened
        """
        # Check if we're in a Docker container
        if IN_DOCKER:
            return {
                "success": False,
                "error": "Browser opening is not available when running in a Docker container.",
                "hint": "Connect from the host machine at the exposed port (usually http://localhost:6789)",
                "url": None,
            }

        try:
            config_dir = default_config_dir()
            config = load_config(config_dir)
            real_port = config["real_port"]
            # Build URL
            url = f"http://localhost:{real_port}"
            if node_id:
                url += f"#{node_id}"

            # Open browser
            webbrowser.open(url)

            logger.info(f"Opened browser to {url}")

            message = "Browser opened to knowledge graph visualization"
            if node_id:
                message += f" focused on node {node_id}"

            return {"success": True, "url": url, "message": message}

        except Exception as e:
            logger.error(f"Failed to open browser: {e}")
            return {"success": False, "error": str(e), "url": None}

    # Vector search tools
    db_path = str(config["database_path"])

    @mcp.tool
    async def search_entities_semantic(
        query: str,
        k: int = 10,
        threshold: float = 0.0,
        model_name: str | None = None,
    ) -> dict[str, Any]:
        """
        Search for entities using semantic similarity via vector embeddings.

        Finds entities whose observations are semantically similar to the query,
        even if they don't share exact keywords. Requires embeddings to be
        generated first.

        Args:
            query (str): Natural language search query
            k (int): Maximum number of results to return (default: 10)
            threshold (float): Minimum similarity score (0-1, default: 0.0)
            model_name (str, optional): Filter by embedding model

        Returns:
            dict: Search results with entity_id, similarity score, and entity data
        """
        from memgraph.embeddings import get_embedding_provider
        from memgraph.vector_sqlite import VectorSQLiteStore

        logger.info(f"Semantic search: {query} (k={k}, threshold={threshold})")

        # Get embedding provider
        provider = get_embedding_provider()
        if provider is None:
            return {
                "error": "No embedding provider configured",
                "hint": "Call configure_embeddings first or set environment variables"
            }

        try:
            # Use context manager for vector store
            with VectorSQLiteStore(db_path=db_path) as vector_store:
                # Generate query embedding
                query_embedding = provider.generate(query)

                # Search vector store
                results = vector_store.search(
                    query_embedding=query_embedding,
                    k=k,
                    threshold=threshold,
                    model_name=model_name,
                )

                # Fetch entity data for results
                entities = []
                for entity_id, score in results:
                    # Use existing search to get entity data
                    entity_data = await DB.search_nodes(entity_id)
                    entities_list = entity_data.get("entities") or []
                    if not entities_list:
                        continue

                    # Prefer an exact name match to the entity_id; fall back to first
                    exact_match = next(
                        (e for e in entities_list if e.get("name") == entity_id),
                        None,
                    )
                    entity_info = dict(exact_match or entities_list[0])
                    entity_info["similarity_score"] = score
                    entities.append(entity_info)

                return {
                    "query": query,
                    "count": len(entities),
                    "results": entities,
                }

        except Exception as e:
            logger.error(f"Semantic search failed: {e}")
            return {"error": str(e)}

    @mcp.tool
    async def search_hybrid(
        query: str,
        k: int = 10,
        semantic_weight: float = 0.7,
        threshold: float = 0.0,
    ) -> dict[str, Any]:
        """
        Hybrid search combining keyword and semantic similarity.

        Merges results from traditional keyword search and vector similarity
        search, with configurable weighting between the two approaches.

        Args:
            query (str): Search query
            k (int): Maximum results to return (default: 10)
            semantic_weight (float): Weight for semantic results (0-1, default: 0.7)
            threshold (float): Minimum similarity threshold (default: 0.0)

        Returns:
            dict: Combined search results ranked by hybrid score
        """
        logger.info(f"Hybrid search: {query} (semantic_weight={semantic_weight})")

        # Get both keyword and semantic results
        keyword_results: dict[str, Any] = await search_nodes(query)  # type: ignore[operator]
        semantic_results: dict[str, Any] = await search_entities_semantic(  # type: ignore[operator]
            query, k=k * 2, threshold=threshold
        )

        # Handle errors
        if "error" in semantic_results:
            logger.warning(f"Semantic search failed, falling back to keyword only: {semantic_results['error']}")
            return keyword_results

        # Build entity scores
        entity_scores: dict[str, dict[str, Any]] = {}

        # Process keyword results (give them keyword_weight score)
        keyword_weight = 1.0 - semantic_weight
        for entity in keyword_results.get("entities", []):
            entity_id = entity.get("name", "")
            entity_scores[entity_id] = {
                "entity": entity,
                "keyword_score": keyword_weight,
                "semantic_score": 0.0,
            }

        # Process semantic results
        for entity in semantic_results.get("results", []):
            entity_id = entity.get("name", "")
            sim_score = entity.pop("similarity_score", 0.0)

            if entity_id in entity_scores:
                # Entity found by both methods
                entity_scores[entity_id]["semantic_score"] = sim_score * semantic_weight
            else:
                # Entity only found by semantic search
                entity_scores[entity_id] = {
                    "entity": entity,
                    "keyword_score": 0.0,
                    "semantic_score": sim_score * semantic_weight,
                }

        # Calculate hybrid scores and sort
        ranked_results = []
        for _entity_id, scores in entity_scores.items():
            hybrid_score = scores["keyword_score"] + scores["semantic_score"]
            result = scores["entity"].copy()
            result["hybrid_score"] = hybrid_score
            result["keyword_score"] = scores["keyword_score"]
            result["semantic_score"] = scores["semantic_score"]
            ranked_results.append(result)

        # Sort by hybrid score descending
        ranked_results.sort(key=lambda x: x["hybrid_score"], reverse=True)

        # Limit to k results
        ranked_results = ranked_results[:k]

        return {
            "query": query,
            "count": len(ranked_results),
            "semantic_weight": semantic_weight,
            "results": ranked_results,
        }

    @mcp.tool
    async def generate_embeddings(
        batch_size: int = 100,
        model_name: str | None = None,
    ) -> dict[str, Any]:
        """
        Generate vector embeddings for all entities that don't have them yet.

        Processes entities in batches to create embeddings from their observations.
        Embeddings enable semantic search capabilities.

        Args:
            batch_size (int): Number of entities to process per batch (default: 100)
            model_name (str, optional): Specific embedding model to use

        Returns:
            dict: Statistics about embeddings generated
        """
        from memgraph.embeddings import get_embedding_provider
        from memgraph.vector_sqlite import VectorSQLiteStore

        logger.info(f"Generating embeddings (batch_size={batch_size})")

        # Get embedding provider
        provider = get_embedding_provider()
        if provider is None:
            return {
                "error": "No embedding provider configured",
                "hint": "Call configure_embeddings first or set environment variables"
            }

        # Initialize vector store
        with VectorSQLiteStore(db_path=db_path) as vector_store:

            try:
                # Get all entities
                graph = await DB.read_graph()
                entities = graph.get("entities", [])

                # Filter entities that need embeddings
                existing_count = vector_store.count(model_name=provider.model_name)
                logger.info(f"Found {len(entities)} total entities, {existing_count} already have embeddings")

                # Get entities without embeddings for this model
                entities_to_process = []
                for entity in entities:
                    entity_id = entity.get("name", "")
                    if not vector_store.exists(entity_id, model_name=provider.model_name):
                        entities_to_process.append(entity)

                if not entities_to_process:
                    return {
                        "message": "All entities already have embeddings",
                        "total_entities": len(entities),
                        "existing_embeddings": existing_count,
                        "generated": 0,
                    }

                # Process in batches
                generated = 0
                for i in range(0, len(entities_to_process), batch_size):
                    batch = entities_to_process[i:i + batch_size]

                    # Create text from observations
                    texts = []
                    entity_ids = []
                    for entity in batch:
                        entity_id = entity.get("name", "")
                        observations = entity.get("observations", [])
                        text = " ".join(observations) if observations else entity_id
                        texts.append(text)
                        entity_ids.append(entity_id)

                    # Generate embeddings
                    embeddings = provider.batch_generate(texts)

                    # Store embeddings
                    vector_store.batch_add(
                        entity_ids=entity_ids,
                        embeddings=embeddings,
                        model_name=provider.model_name,
                    )

                    generated += len(batch)
                    logger.info(f"Generated {generated}/{len(entities_to_process)} embeddings")

                return {
                    "message": f"Generated {generated} embeddings",
                    "total_entities": len(entities),
                    "existing_embeddings": existing_count,
                    "generated": generated,
                    "model": provider.model_name,
                }

            except Exception as e:
                logger.error(f"Failed to generate embeddings: {e}")
                return {"error": str(e)}

    @mcp.tool
    async def configure_embeddings(
        provider: str = "sentence-transformers",
        model: str | None = None,
        api_key: str | None = None,
    ) -> dict[str, Any]:
        """
        Configure the embedding provider for semantic search.

        Supports local sentence-transformers models or OpenAI's API.

        Args:
            provider (str): "sentence-transformers" or "openai" (default: sentence-transformers)
            model (str, optional): Model name. Defaults:
                - sentence-transformers: "all-MiniLM-L6-v2"
                - openai: "text-embedding-3-small"
            api_key (str, optional): OpenAI API key (for OpenAI provider)

        Returns:
            dict: Configuration status and provider details
        """
        from memgraph.embeddings import configure_from_dict, get_embedding_provider

        logger.info(f"Configuring embeddings: provider={provider}, model={model}")

        try:
            config_dict = {
                "provider": provider,
                "model": model,
                "api_key": api_key,
            }

            configure_from_dict(config_dict)

            # Verify configuration
            active_provider = get_embedding_provider()
            if active_provider is None:
                return {"error": "Failed to configure provider"}

            return {
                "message": "Embedding provider configured successfully",
                "provider": provider,
                "model": active_provider.model_name,
                "dimensions": active_provider.dimensions,
            }

        except Exception as e:
            logger.error(f"Failed to configure embeddings: {e}")
            return {"error": str(e)}

    return mcp


def get_lifespan_hook(config: Config) -> Lifespan[Any]:
    """
    Create an async lifespan hook for the FastMCP application.
    """

    @asynccontextmanager
    async def lifecycle_hook(app: FastAPI) -> AsyncGenerator[None, Any]:
        """Example of an async lifecycle hook for the unified app."""

        info_file = save_server_info(
            config["config_dir"],
            launched_by="unified_service",
            pid=os.getpid(),
            host=config["host"],
            port=config["port"],
            database_path=config["database_path"],
        )

        def cleanup() -> None:
            if info_file:
                info_file.unlink(missing_ok=True)

        atexit.register(cleanup)
        try:
            yield
        finally:
            info_file.unlink(missing_ok=True)

    return lifecycle_hook


if __name__ == "__main__":
    # Run the MCP server
    import sys

    try:
        mcp = setup_mcp(load_config(default_config_dir()))
        mcp.run()
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Server error: {e}")
        sys.exit(1)


# Export tool functions for testing
# These are wrappers that delegate to the module-level mcp instance
async def configure_embeddings(
    provider: str = "sentence-transformers",
    model: str | None = None,
    api_key: str | None = None,
) -> dict[str, Any]:
    """Wrapper for testing - calls the mcp tool."""
    mcp = get_mcp()
    result = await mcp.call_tool("configure_embeddings", {"provider": provider, "model": model, "api_key": api_key})
    return result.structured_content  # type: ignore[return-value]


async def generate_embeddings(
    batch_size: int = 100,
    model_name: str | None = None,
) -> dict[str, Any]:
    """Wrapper for testing - calls the mcp tool."""
    mcp = get_mcp()
    result = await mcp.call_tool("generate_embeddings", {"batch_size": batch_size, "model_name": model_name})
    return result.structured_content  # type: ignore[return-value]


async def search_entities_semantic(
    query: str,
    k: int = 10,
    threshold: float = 0.0,
    model_name: str | None = None,
) -> dict[str, Any]:
    """Wrapper for testing - calls the mcp tool."""
    mcp = get_mcp()
    result = await mcp.call_tool(
        "search_entities_semantic",
        {"query": query, "k": k, "threshold": threshold, "model_name": model_name},
    )
    return result.structured_content  # type: ignore[return-value]


async def search_hybrid(
    query: str,
    k: int = 10,
    semantic_weight: float = 0.7,
    threshold: float = 0.0,
) -> dict[str, Any]:
    """Wrapper for testing - calls the mcp tool."""
    mcp = get_mcp()
    result = await mcp.call_tool(
        "search_hybrid",
        {"query": query, "k": k, "semantic_weight": semantic_weight, "threshold": threshold},
    )
    return result.structured_content  # type: ignore[return-value]
