"""
Genesis V2 Native MCP Server
==============================
Core tool implementations for the Genesis MCP server.
Provides: KG tools, browser tools, identity, health check, configure.

Reconstructed from compiled bytecode (native_server.cpython-312.pyc).
"""

import asyncio
import json
import logging
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Any, Optional

from genesis_v2.core.kg_adapter import create_adapter, KGAdapter
from genesis_v2.core.identity import Identity
from genesis_v2.core.browser import get_browser_controller
from genesis_v2.mcp.validation import (
    validate_url,
    sanitize_selector,
    validate_kg_search_input,
    validate_entity_json,
    check_rate_limit,
)
from genesis_v2.mcp.antigravity_tools import (
    antigravity_plan,
    antigravity_execute,
    gemini_query,
    sync_push,
    sync_pull,
    sync_status,
    memory_query,
    memory_store,
    antigravity_browser,
)

logger = logging.getLogger(__name__)


# =============================================================================
# Entity wrapper — gives dicts a .to_dict() method
# =============================================================================

class Entity:
    """Wraps a KG entity dict to provide object-style access."""

    def __init__(self, data: dict):
        self._data = data

    @property
    def id(self) -> str:
        return self._data.get("id", "")

    @property
    def type(self) -> str:
        return self._data.get("type", "")

    def to_dict(self) -> dict:
        return self._data

    @classmethod
    def from_dict(cls, data: dict) -> "Entity":
        return cls(data)

    def __repr__(self):
        return f"Entity(id={self.id!r}, type={self.type!r})"


# =============================================================================
# Extended KGAdapter — adds methods expected by native_server
# =============================================================================

class ExtendedKGAdapter:
    """
    Extends the base KGAdapter with the richer API expected by native_server:
    connect(), count(), list_types(), get(), ingest(), get_context()
    """

    def __init__(self, base: KGAdapter):
        self._base = base

    def connect(self):
        """Ensure KG data is loaded (connection analogue)."""
        if not self._base._loaded:
            self._base.load()

    def count(self) -> int:
        """Return total entity count."""
        if not self._base._loaded:
            self._base.load()
        return len(self._base.entities)

    def list_types(self) -> List[str]:
        """Return list of distinct entity types."""
        if not self._base._loaded:
            self._base.load()
        return sorted({e.get("type", "unknown") for e in self._base.entities})

    def search(self, query: str, limit: int = 10) -> List[Entity]:
        """Search and return Entity objects."""
        results = self._base.search(query, limit)
        return [Entity(r) for r in results]

    def get(self, entity_id: str) -> Optional[Entity]:
        """Get entity by ID."""
        if not self._base._loaded:
            self._base.load()
        for e in self._base.entities:
            if e.get("id") == entity_id:
                return Entity(e)
        return None

    def ingest(self, entity: Entity) -> bool:
        """Add entity to the in-memory KG (and optionally persist)."""
        if not self._base._loaded:
            self._base.load()
        # Replace existing or append
        for i, e in enumerate(self._base.entities):
            if e.get("id") == entity.id:
                self._base.entities[i] = entity.to_dict()
                return True
        self._base.entities.append(entity.to_dict())
        return True

    def get_context(self, max_chars: int = 2000) -> str:
        """Return a summary string suitable for agent injection."""
        if not self._base._loaded:
            self._base.load()
        total = self.count()
        types = self.list_types()
        summary = f"Knowledge Graph: {total} entities across {len(types)} types.\nTypes: {', '.join(types[:20])}"
        if len(summary) > max_chars:
            summary = summary[:max_chars]
        return summary


# =============================================================================
# Server context — immutable container for all dependencies
# =============================================================================

@dataclass
class ServerContext:
    """
    Immutable server context with all dependencies.
    Replaces global mutable state for thread safety.
    """
    kg_adapter: Optional[ExtendedKGAdapter] = None
    identity: Optional[Identity] = None

    @classmethod
    def create(
        cls,
        kg_path: Optional[str] = None,
        identity_rules_path: Optional[str] = None,
    ) -> "ServerContext":
        """Factory method to create configured context."""
        kg_adapter = None

        if kg_path:
            base = create_adapter(kg_path)
            kg_adapter = ExtendedKGAdapter(base)
            kg_adapter.connect()
            logger.info(f"KG connected: {kg_adapter.count()} entities")
        else:
            # Try default path
            default_kg = Path(__file__).parent.parent.parent / "KNOWLEDGE_GRAPH" / "entities.jsonl"
            if default_kg.exists():
                base = create_adapter(str(default_kg))
                kg_adapter = ExtendedKGAdapter(base)
                kg_adapter.connect()
                logger.info(f"KG connected (default): {kg_adapter.count()} entities")

        identity = Identity()
        identity.load()
        logger.info(f"Identity loaded: {identity.name}")

        return cls(kg_adapter=kg_adapter, identity=identity)


# Global server context (lazy-initialized)
_context: Optional[ServerContext] = None


def _get_context() -> ServerContext:
    """Get or create the server context."""
    global _context
    if _context is None:
        _context = ServerContext.create()
    return _context


# =============================================================================
# KG tools
# =============================================================================

def kg_search(query: str, limit: int = 10) -> str:
    """
    Search the knowledge graph for entities matching query.

    Args:
        query: Search query string
        limit: Maximum results to return (1-100)

    Returns:
        JSON string with matching entities
    """
    allowed, error = check_rate_limit("kg_search")
    if not allowed:
        return json.dumps({"error": error})

    valid, error = validate_kg_search_input(query, limit)
    if not valid:
        return json.dumps({"error": error})

    ctx = _get_context()
    if ctx.kg_adapter is None:
        return json.dumps({"error": "Knowledge graph not connected"})

    results = ctx.kg_adapter.search(query, limit)
    return json.dumps(
        {"query": query, "count": len(results), "entities": [e.to_dict() for e in results]},
        indent=2,
    )


def kg_get(entity_id: str) -> str:
    """
    Get a specific entity by ID.

    Args:
        entity_id: The entity ID to retrieve

    Returns:
        JSON string with entity data or error
    """
    ctx = _get_context()
    if ctx.kg_adapter is None:
        return json.dumps({"error": "Knowledge graph not connected"})

    if not entity_id or not isinstance(entity_id, str):
        return json.dumps({"error": "entity_id must be a non-empty string"})

    entity = ctx.kg_adapter.get(entity_id.strip())
    if entity:
        return json.dumps(entity.to_dict(), indent=2)
    return json.dumps({"error": f"Entity not found: {entity_id}"})


def kg_ingest(entity_json: str) -> str:
    """
    Add new entity to knowledge graph.

    Args:
        entity_json: JSON string with entity data (must have 'id' and 'type')

    Returns:
        Success/failure message
    """
    ctx = _get_context()
    if ctx.kg_adapter is None:
        return json.dumps({"error": "Knowledge graph not connected"})

    valid, error, data = validate_entity_json(entity_json)
    if not valid:
        return json.dumps({"error": error})

    try:
        entity = Entity.from_dict(data)
        success = ctx.kg_adapter.ingest(entity)
        return json.dumps({"success": success, "entity_id": entity.id})
    except Exception as e:
        return json.dumps({"error": str(e)})


def kg_context() -> str:
    """
    Get knowledge graph context summary for agent injection.

    Returns:
        Summary string suitable for system instruction
    """
    ctx = _get_context()
    if ctx.kg_adapter is None:
        return "Knowledge graph not connected"
    return ctx.kg_adapter.get_context(2000)


def kg_stats() -> str:
    """
    Get knowledge graph statistics.

    Returns:
        JSON with entity count and types
    """
    ctx = _get_context()
    if ctx.kg_adapter is None:
        return json.dumps({"error": "Knowledge graph not connected"})
    return json.dumps(
        {
            "connected": True,
            "entity_count": ctx.kg_adapter.count(),
            "entity_types": ctx.kg_adapter.list_types(),
        },
        indent=2,
    )


# =============================================================================
# Browser tools
# =============================================================================

def browser_navigate(url: str) -> str:
    """
    Navigate browser to URL.

    Args:
        url: Target URL (must be http/https)

    Returns:
        Navigation result
    """
    valid, error = validate_url(url)
    if not valid:
        return json.dumps({"success": False, "error": error})

    controller = get_browser_controller()
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(controller.navigate(url))
        return json.dumps({"success": result.success, "url": url, "data": result.data})
    finally:
        loop.close()


def browser_click(selector: str) -> str:
    """
    Click element by CSS selector.

    Args:
        selector: CSS selector (sanitized)

    Returns:
        Click result
    """
    clean_selector, error = sanitize_selector(selector)
    if error:
        return json.dumps({"success": False, "error": error})

    controller = get_browser_controller()
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(controller.click(clean_selector))
        return json.dumps({"success": result.success, "selector": clean_selector, "data": result.data})
    finally:
        loop.close()


def browser_screenshot() -> str:
    """
    Capture browser screenshot.

    Returns:
        Screenshot metadata (base64 data if available)
    """
    controller = get_browser_controller()
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(controller.screenshot())
        return json.dumps({"success": result.success, "data": result.data})
    finally:
        loop.close()


# =============================================================================
# Identity / health tools
# =============================================================================

def get_identity() -> str:
    """
    Get agent identity and persona information.

    Returns:
        Identity JSON with persona name and rules
    """
    ctx = _get_context()
    identity = ctx.identity
    if identity is None:
        identity = Identity()
        identity.load()

    rules_preview = None
    if identity.persona:
        rules_preview = identity.persona[:500] + "..." if len(identity.persona) > 500 else identity.persona

    return json.dumps(
        {
            "persona": identity.name,
            "rules_preview": rules_preview,
        },
        indent=2,
    )


def health_check() -> str:
    """
    Check server health status.

    Returns:
        JSON with health status of all components
    """
    ctx = _get_context()
    checks = {"status": "healthy", "components": {}}

    # KG health
    try:
        if ctx.kg_adapter:
            count = ctx.kg_adapter.count()
            checks["components"]["knowledge_graph"] = {"status": "healthy", "entity_count": count}
        else:
            checks["components"]["knowledge_graph"] = {"status": "not_configured"}
    except Exception as e:
        checks["components"]["knowledge_graph"] = {"status": "unhealthy", "error": str(e)}
        checks["status"] = "degraded"

    # Identity health
    try:
        if ctx.identity and (ctx.identity.persona or ctx.identity.name):
            checks["components"]["identity"] = {"status": "healthy", "persona": ctx.identity.name}
        else:
            checks["components"]["identity"] = {"status": "not_configured"}
    except Exception as e:
        checks["components"]["identity"] = {"status": "unhealthy", "error": str(e)}
        checks["status"] = "degraded"

    # Browser health
    try:
        controller = get_browser_controller()
        checks["components"]["browser"] = {"status": "healthy", "mode": "stub"}
    except Exception as e:
        checks["components"]["browser"] = {"status": "unhealthy", "error": str(e)}
        checks["status"] = "degraded"

    return json.dumps(checks, indent=2)


# =============================================================================
# Configuration / utility
# =============================================================================

def get_mcp_tools():
    """Get all tools as ADK FunctionTools."""
    try:
        from google.adk.tools import FunctionTool
        from genesis_v2.core.adk_orchestrator import adk_execute_browser_action, adk_extract_page_content
        return [
            FunctionTool(func=kg_search),
            FunctionTool(func=kg_get),
            FunctionTool(func=kg_ingest),
            FunctionTool(func=kg_context),
            FunctionTool(func=kg_stats),
            FunctionTool(func=browser_navigate),
            FunctionTool(func=browser_click),
            FunctionTool(func=browser_screenshot),
            FunctionTool(func=get_identity),
            FunctionTool(func=health_check),
            FunctionTool(func=antigravity_plan),
            FunctionTool(func=antigravity_execute),
            FunctionTool(func=gemini_query),
            FunctionTool(func=sync_push),
            FunctionTool(func=sync_pull),
            FunctionTool(func=sync_status),
            FunctionTool(func=memory_query),
            FunctionTool(func=memory_store),
            FunctionTool(func=antigravity_browser),
            FunctionTool(func=adk_execute_browser_action),
            FunctionTool(func=adk_extract_page_content),
        ]
    except ImportError:
        logger.warning("google.adk not available — get_mcp_tools() returning empty list")
        return []


def configure(
    kg_path: Optional[str] = None,
    identity_rules_path: Optional[str] = None,
):
    """
    Configure the MCP server with knowledge graph and identity.

    Args:
        kg_path: Path to knowledge graph (JSONL file)
        identity_rules_path: Path to identity rules file
    """
    global _context
    _context = ServerContext.create(kg_path, identity_rules_path)


def test_server() -> dict:
    """Test server without starting it."""
    configure()
    ctx = _get_context()
    tools = get_mcp_tools()
    return {
        "status": "ready",
        "kg_connected": ctx.kg_adapter is not None,
        "kg_entities": ctx.kg_adapter.count() if ctx.kg_adapter else 0,
        "identity": ctx.identity.name if ctx.identity else None,
        "tools": [t.func.__name__ for t in tools],
        "tool_count": len(tools),
    }


# =============================================================================
# Standalone test entry point
# =============================================================================

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    if len(sys.argv) > 1 and sys.argv[1] == "test":
        result = test_server()
        for key, value in result.items():
            print(f"{key}: {value}")