#!/usr/bin/env python3
"""
Siemens Graph Studio MCP Server

This server provides intelligent SPARQL query capabilities for Siemens Graph Studio
(formerly Altair Graph Studio / AGS). It exposes knowledge graph operations via the
Model Context Protocol (MCP), enabling AI agents to interact with GraphMarts.

For usage information, run:
    siemens-graph-studio-mcp --help

Or see the README.md file.
"""

import os
import sys
import json
import argparse
import logging

# Add package directory to path for imports
sys.path.insert(0, os.path.dirname(__file__))

try:
    from mcp.server.fastmcp import FastMCP
    from .models import (
        GraphmartConfig,
        QuestionResponse,
        KnowledgeExploration,
        ExamplesResponse,
        AgentInsights,
        QueryResult
    )
    from .tools import register_tools, get_tool_categories
    
    try:
        from .sparql_agent_core import SPARQLAgent
    except ImportError:
        print("⚠️  Warning: SPARQLAgent core not available - using stub", file=sys.stderr)
        SPARQLAgent = None
        
except ImportError as e:
    print(f"❌ Error importing dependencies: {e}", file=sys.stderr)
    print("   Please install with: pip install siemens-graph-studio-mcp-server", file=sys.stderr)
    sys.exit(1)

# Configure logging to stderr to avoid interfering with MCP JSON-RPC on stdout
logging.basicConfig(
    level=logging.ERROR,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]
)
logger = logging.getLogger(__name__)

# Global variables
mcp = None
agent = None

# Mapping from config file keys to environment variable names
_CONFIG_ENV_MAP = {
    "anzo_server":    "ANZO_SERVER",
    "anzo_port":      "ANZO_PORT",
    "anzo_username":  "ANZO_USERNAME",
    "anzo_password":  "ANZO_PASSWORD",
    "graphmart_uri":  "GRAPHMART_URI",
    "api_key":        "API_KEY",
    "enable_agent_debug":   "ENABLE_AGENT_DEBUG",
    "enable_logging_debug": "ENABLE_LOGGING_DEBUG",
}


def _strip_jsonc_comments(text: str) -> str:
    """Strip // line comments and /* */ block comments from a JSON string."""
    import re
    text = re.sub(r'/\*.*?\*/', '', text, flags=re.DOTALL)
    text = re.sub(r'(?<!:)//[^\n]*', '', text)
    return text


def load_config_file(config_path: str) -> dict:
    """Load a config file (JSON/JSONC or .env) and inject connection
    settings into os.environ.

    Supported formats:
        - .json / .jsonc: JSON with optional // comments
        - .env: KEY=VALUE format (one per line)

    Config file keys (all optional):
        anzo_server, anzo_port, anzo_username, anzo_password,
        graphmart_uri, api_key, enable_agent_debug, enable_logging_debug,
        mcp_transport, mcp_host, mcp_port

    Returns the raw config dict so callers can read transport settings.
    """
    if not os.path.exists(config_path):
        print(f"❌ Config file not found: {config_path}", file=sys.stderr)
        sys.exit(1)

    ext = os.path.splitext(config_path)[1].lower()

    if ext == '.env':
        return _load_env_config(config_path)
    elif ext in ('.json', '.jsonc'):
        return _load_json_config(config_path)
    else:
        print(f"❌ Unsupported config format '{ext}'. Use .json, .jsonc, or .env", file=sys.stderr)
        print("   Examples:", file=sys.stderr)
        print("     siemens-graph-studio-mcp --config connection.json", file=sys.stderr)
        print("     siemens-graph-studio-mcp --config .env", file=sys.stderr)
        sys.exit(1)


def _load_env_config(config_path: str) -> dict:
    """Load a .env file and inject values into os.environ."""
    cfg = {}
    # Build reverse map: ENV_VAR -> config_key
    env_to_key = {v: k for k, v in _CONFIG_ENV_MAP.items()}
    
    with open(config_path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            if '=' not in line:
                continue
            key, _, value = line.partition('=')
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            os.environ[key] = value
            
            # Track in cfg dict for transport settings
            config_key = env_to_key.get(key)
            if config_key:
                cfg[config_key] = value
    
    print(f"🔧 Loaded config from {config_path} (.env format)", file=sys.stderr)
    return cfg


def _load_json_config(config_path: str) -> dict:
    """Load a JSON/JSONC config file and inject into os.environ."""
    with open(config_path) as f:
        raw = f.read()

    try:
        cfg = json.loads(_strip_jsonc_comments(raw))
    except json.JSONDecodeError as e:
        print(f"❌ Failed to parse config file {config_path}: {e}", file=sys.stderr)
        sys.exit(1)

    injected = []
    for key, env_var in _CONFIG_ENV_MAP.items():
        if key in cfg:
            os.environ[env_var] = str(cfg[key])
            # Mask secrets in the log
            display = "***" if "password" in key or "api_key" in key else str(cfg[key])
            injected.append(f"{env_var}={display}")

    print(f"🔧 Loaded config from {config_path}", file=sys.stderr)
    for item in injected:
        print(f"   {item}", file=sys.stderr)

    return cfg


def reset_agent():
    """Reset the global agent instance."""
    global agent
    agent = None


async def get_agent():
    """Get the SPARQL agent instance - fails gracefully if environment not configured."""
    global agent
    if agent is None:
        if not SPARQLAgent:
            raise RuntimeError("SPARQLAgent core not available - check imports")
        
        try:
            # All configuration comes from environment variables - no hardcoded defaults
            anzo_server = os.environ.get('ANZO_SERVER')
            anzo_port = os.environ.get('ANZO_PORT')
            anzo_username = os.environ.get('ANZO_USERNAME')
            anzo_password = os.environ.get('ANZO_PASSWORD')
            graphmart_uri = os.environ.get('GRAPHMART_URI')
            
            # Validate required environment variables
            missing = []
            if not anzo_server:
                missing.append('ANZO_SERVER')
            if not anzo_port:
                missing.append('ANZO_PORT')
            if not anzo_username:
                missing.append('ANZO_USERNAME')
            if not anzo_password:
                missing.append('ANZO_PASSWORD')
            if not graphmart_uri:
                missing.append('GRAPHMART_URI')
                
            if missing:
                raise RuntimeError(
                    f"Missing required environment variables: {', '.join(missing)}. "
                    "See README.md for configuration instructions."
                )
            
            config = GraphmartConfig(
                ags_server=anzo_server,
                ags_port=int(anzo_port),
                graphmart_uri=graphmart_uri,
                username=anzo_username,
                password=anzo_password
            )
            
            logger.info(f"Initializing Graph Studio agent with server: {config.ags_server}:{config.ags_port}")
            
            agent = SPARQLAgent(config)
            await agent.initialize()
            logger.info("SPARQL Agent initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize agent: {e}")
            raise RuntimeError(
                f"Failed to initialize Graph Studio agent: {e}. "
                "Check environment variables: ANZO_SERVER, ANZO_PORT, "
                "ANZO_USERNAME, ANZO_PASSWORD, GRAPHMART_URI, API_KEY"
            )
    
    return agent


def run_verification():
    """Run installation verification and print diagnostic information."""
    import shutil
    import platform
    
    print("=" * 60)
    print("Siemens Graph Studio MCP Server - Installation Verification")
    print("=" * 60)
    print()
    
    all_ok = True
    
    # 1. Python version check
    print("🐍 Python Environment")
    print(f"   Python version: {platform.python_version()}")
    print(f"   Python path: {sys.executable}")
    
    if sys.version_info < (3, 10):
        print("   ❌ ERROR: Python 3.10+ required")
        all_ok = False
    else:
        print("   ✅ Python version OK")
    print()
    
    # 2. Package installation check
    print("📦 Package Installation")
    try:
        import siemens_graph_studio_mcp
        pkg_path = os.path.dirname(siemens_graph_studio_mcp.__file__)
        print(f"   Package location: {pkg_path}")
        print("   ✅ Package installed correctly")
    except ImportError as e:
        print(f"   ❌ ERROR: Package not found - {e}")
        all_ok = False
    print()
    
    # 3. Dependencies check
    print("📚 Dependencies")
    deps = [
        ("mcp", "mcp"),
        ("pydantic", "pydantic"),
        ("pyanzo", "pyanzo"),
        ("openai", "openai"),
        ("rdflib", "rdflib"),
    ]
    for name, module in deps:
        try:
            __import__(module)
            print(f"   ✅ {name}")
        except ImportError:
            print(f"   ❌ {name} - NOT INSTALLED")
            all_ok = False
    print()
    
    # 4. Command availability
    print("🔧 Command Line Tool")
    cmd_path = shutil.which("siemens-graph-studio-mcp")
    if cmd_path:
        print(f"   Command path: {cmd_path}")
        print("   ✅ Command available in PATH")
    else:
        print("   ⚠️  Command not found in PATH")
        print("   Try running directly with Python:")
        print(f"   {sys.executable} -m siemens_graph_studio_mcp.server")
        # Find potential installation location
        user_base = os.path.join(os.path.expanduser("~"), ".local", "bin")
        print(f"\n   Or add to PATH: export PATH=\"{user_base}:$PATH\"")
    print()
    
    # 5. Environment variables
    print("🔐 Environment Variables")
    env_vars = {
        "ANZO_SERVER": True,      # required
        "ANZO_PORT": True,        # required
        "ANZO_USERNAME": True,    # required
        "ANZO_PASSWORD": True,    # required (mask value)
        "GRAPHMART_URI": True,    # required
        "API_KEY": False,         # optional
    }
    
    missing_required = []
    for var, required in env_vars.items():
        value = os.environ.get(var)
        if value:
            # Mask sensitive values
            if "PASSWORD" in var or "KEY" in var or "SECRET" in var:
                display = "***" + value[-4:] if len(value) > 4 else "***"
            else:
                display = value[:40] + "..." if len(value) > 40 else value
            print(f"   ✅ {var}={display}")
        else:
            if required:
                print(f"   ❌ {var} - NOT SET (required)")
                missing_required.append(var)
            else:
                print(f"   ⚪ {var} - not set (optional)")
    
    if missing_required:
        all_ok = False
        print(f"\n   ⚠️  Set required variables before using the MCP server")
    print()
    
    # 6. VS Code MCP configuration
    print("🔌 VS Code MCP Configuration")
    vscode_mcp_paths = [
        os.path.join(os.getcwd(), ".vscode", "mcp.json"),
        os.path.join(os.path.expanduser("~"), ".vscode", "mcp.json"),
    ]
    mcp_found = False
    for path in vscode_mcp_paths:
        if os.path.exists(path):
            print(f"   Found config: {path}")
            mcp_found = True
    
    if mcp_found:
        print("   ✅ MCP configuration found")
    else:
        print("   ⚪ No .vscode/mcp.json found in current directory")
        print("   See docs/INSTALLATION_GUIDE.md for setup instructions")
    print()
    
    # Summary
    print("=" * 60)
    if all_ok:
        print("✅ All checks passed! Installation looks good.")
        print("\nNext steps:")
        print("1. Configure .vscode/mcp.json (see docs/INSTALLATION_GUIDE.md)")
        print("2. Enable MCP in VS Code settings")
        print("3. Reload VS Code window")
    else:
        print("❌ Some issues found. See above for details.")
        print("\nFor troubleshooting, see: docs/INSTALLATION_GUIDE.md")
    print("=" * 60)


def setup_vscode_config():
    """Generate .vscode/mcp.json in the current directory."""
    import shutil
    
    vscode_dir = os.path.join(os.getcwd(), ".vscode")
    mcp_json_path = os.path.join(vscode_dir, "mcp.json")
    
    # Check if file already exists
    if os.path.exists(mcp_json_path):
        print(f"⚠️  {mcp_json_path} already exists.")
        response = input("   Overwrite? [y/N]: ").strip().lower()
        if response != 'y':
            print("   Aborted.")
            return
    
    # Find the command path
    cmd_path = shutil.which("siemens-graph-studio-mcp")
    if not cmd_path:
        cmd_path = "siemens-graph-studio-mcp"
        print("⚠️  Command not found in PATH - using name only.", file=sys.stderr)
        print("   You may need to replace with the full path in mcp.json", file=sys.stderr)
    
    config = {
        "servers": {
            "graph-studio": {
                "type": "stdio",
                "command": cmd_path,
                "env": {
                    "ANZO_SERVER": "your-server.example.com",
                    "ANZO_PORT": "8443",
                    "ANZO_USERNAME": "your-username",
                    "ANZO_PASSWORD": "your-password",
                    "GRAPHMART_URI": "http://example.com/Graphmart/your-graphmart-id"
                }
            }
        }
    }
    
    os.makedirs(vscode_dir, exist_ok=True)
    with open(mcp_json_path, 'w') as f:
        json.dump(config, f, indent=2)
        f.write('\n')
    
    print(f"✅ Created {mcp_json_path}")
    print()
    print("Next steps:")
    print("  1. Edit .vscode/mcp.json with your Graph Studio credentials")
    print("  2. Reload VS Code (Cmd+Shift+P → 'Developer: Reload Window')")
    print("  3. Open Copilot Chat (Cmd+Shift+I) and start querying!")
    print()
    print("📖 See docs/INSTALLATION_GUIDE.md for full setup instructions")


def main():
    """Main entry point for the MCP server."""
    global mcp

    parser = argparse.ArgumentParser(
        description="Siemens Graph Studio MCP Server - AI agent interface for knowledge graphs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Environment Variables:
  ANZO_SERVER       Graph Studio server hostname (required)
  ANZO_PORT         Graph Studio server port (required)
  ANZO_USERNAME     Authentication username (required)
  ANZO_PASSWORD     Authentication password (required)
  GRAPHMART_URI     Target GraphMart URI (required)
  API_KEY           OpenAI API key (optional - only for NL→SPARQL features)
                    Basic SPARQL operations work without API_KEY

Examples:
  # Run with environment variables
  export ANZO_SERVER=graph-studio.example.com
  export ANZO_PORT=8443
  export ANZO_USERNAME=admin
  export ANZO_PASSWORD=secret
  export GRAPHMART_URI=http://example.com/Graphmart/my-graphmart
  siemens-graph-studio-mcp

  # Run with config file
  siemens-graph-studio-mcp --config connection.json

  # Run with HTTP transport
  siemens-graph-studio-mcp --transport sse --port 8000
        """
    )
    parser.add_argument(
        "--config", "-c",
        metavar="FILE",
        help="Path to a JSON/JSONC config file with connection settings"
    )
    parser.add_argument(
        "--transport", "-t",
        default=None,
        choices=["stdio", "sse", "streamable-http"],
        help="MCP transport mode (default: stdio)"
    )
    parser.add_argument(
        "--host",
        default=None,
        help="Bind host for HTTP transports (default: 0.0.0.0)"
    )
    parser.add_argument(
        "--port", "-p",
        type=int,
        default=None,
        help="Bind port for HTTP transports (default: 8000)"
    )
    parser.add_argument(
        "--version", "-v",
        action="version",
        version="%(prog)s 0.3.2"
    )
    parser.add_argument(
        "--verify",
        action="store_true",
        help="Verify installation and print diagnostic information"
    )
    parser.add_argument(
        "--setup-vscode",
        action="store_true",
        help="Generate .vscode/mcp.json config in the current directory"
    )
    args = parser.parse_args()
    
    # Handle --verify flag
    if args.verify:
        run_verification()
        sys.exit(0)
    
    # Handle --setup-vscode flag
    if args.setup_vscode:
        setup_vscode_config()
        sys.exit(0)

    # Load config file if provided
    file_cfg = {}
    if args.config:
        file_cfg = load_config_file(args.config)

    # Resolve transport/host/port with priority: CLI arg > config file > default
    transport = args.transport or file_cfg.get("mcp_transport", "stdio")
    host = args.host or file_cfg.get("mcp_host", "0.0.0.0")
    port = args.port or file_cfg.get("mcp_port", 8000)

    # Log startup to stderr
    print("🚀 Starting Siemens Graph Studio MCP Server", file=sys.stderr)
    print(f"   Transport: {transport}", file=sys.stderr)
    if transport != "stdio":
        print(f"   Listening on: {host}:{port}", file=sys.stderr)
    
    # Create FastMCP server instance
    try:
        mcp = FastMCP("Siemens Graph Studio", host=host, port=int(port))
        print("✅ MCP server created", file=sys.stderr)
    except Exception as e:
        print(f"❌ Failed to create MCP server: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Register tools using the modular system
    try:
        tool_count = register_tools(mcp)
        print(f"✅ {tool_count} tools registered", file=sys.stderr)
        
        categories = get_tool_categories()
        for category, tools in categories.items():
            print(f"   📁 {category}: {len(tools)} tools", file=sys.stderr)
            
    except Exception as e:
        print(f"❌ Failed to register tools: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Check required environment variables
    required_env_vars = ['ANZO_SERVER', 'ANZO_PORT', 'ANZO_USERNAME', 'ANZO_PASSWORD', 'GRAPHMART_URI']
    missing_vars = [var for var in required_env_vars if not os.getenv(var)]
    
    if missing_vars:
        print(f"⚠️  Missing environment variables: {missing_vars}", file=sys.stderr)
        print("   Some operations may fail until these are configured", file=sys.stderr)
        print("   Use --config to load from a config file, or set environment variables", file=sys.stderr)
    else:
        print("✅ All required environment variables configured", file=sys.stderr)
        print(f"   Server: {os.getenv('ANZO_SERVER')}:{os.getenv('ANZO_PORT')}", file=sys.stderr)
        print(f"   GraphMart: {os.getenv('GRAPHMART_URI')}", file=sys.stderr)
    
    # Run server
    if transport == "stdio":
        print("📡 Starting MCP server (stdio transport)...", file=sys.stderr)
    else:
        print(f"📡 Starting MCP server ({transport} transport on {host}:{port})...", file=sys.stderr)

    try:
        if transport == "stdio":
            mcp.run()
        else:
            mcp.run(transport=transport)
    except Exception as e:
        print(f"❌ Server error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
