mcp_server.py

#!/usr/bin/env python3
"""
MCP Server for OmniTech Customer Support Chatbot - FULL VERSION
═══════════════════════════════════════════════════════════════════════════════

OVERVIEW
--------
This file implements a Model Context Protocol (MCP) server that provides tools
and resources for the OmniTech Customer Support chatbot. MCP is a protocol that
allows LLM applications to connect to external data sources and tools.

WHAT IS MCP?
------------
MCP (Model Context Protocol) is a standard protocol developed by Anthropic for
connecting LLMs to external tools and data. Think of it like a USB standard for
AI - any MCP-compatible client can connect to any MCP-compatible server.

Key MCP concepts:
- **Tools**: Functions the LLM can call (like classify_query, lookup_customer)
- **Resources**: Data the LLM can read (like config://llm, data://tickets)
- **Server**: Provides tools and resources (this file)
- **Client**: Connects to server and uses tools (rag_agent.py)

This server uses FastMCP — the recommended high-level framework for building
MCP servers. FastMCP auto-generates tool schemas from Python type hints and
docstrings, eliminating the boilerplate of manual Tool() definitions.

WHAT THIS SERVER PROVIDES
-------------------------
1. CLASSIFICATION TOOLS - Categorize customer queries
   - classify_query: Determine which support category a query belongs to
   - get_query_template: Get the prompt template for a category
   - list_categories: List all available support categories

2. KNOWLEDGE TOOLS - Search the vector database
   - search_knowledge: Search for relevant documents
   - get_knowledge_for_query: Get concatenated knowledge for RAG

3. CUSTOMER TOOLS - Manage customer data
   - lookup_customer: Find customer info by email
   - create_support_ticket: Create a new support ticket
   - get_tickets: Retrieve tickets with optional filters

4. STATISTICS TOOLS
   - get_server_stats: Get server health and metrics

5. MCP RESOURCES - Read-only data access
   - config://llm: Current LLM model configuration
   - config://database: Database statistics
   - config://categories: Support category definitions
   - data://tickets: Current support tickets

ARCHITECTURE
------------
    ┌─────────────────────────────────────────────────────────────┐
    │                    rag_agent.py (Client)                    │
    │                    Uses mcp.ClientSession                   │
    └─────────────────────────────────────────────────────────────┘
                                  │
                                  │ stdio (JSON-RPC)
                                  ▼
    ┌─────────────────────────────────────────────────────────────┐
    │                    mcp_server.py (Server)                   │
    │                    FastMCP + OmniTechSupport                │
    ├─────────────────────────────────────────────────────────────┤
    │  Tools: classify_query, search_knowledge, lookup_customer   │
    │  Resources: config://llm, config://database, data://tickets │
    ├─────────────────────────────────────────────────────────────┤
    │  ChromaDB (Vector Store)     │    SQLite (Customer DB)      │
    │  - PDF documents             │    - Customers table         │
    │  - Embeddings                │    - Orders table            │
    │  - Semantic search           │    - Tickets table           │
    └─────────────────────────────────────────────────────────────┘

HOW TO RUN
----------
This server is typically started by rag_agent.py as a subprocess:
    python mcp_server.py

It communicates via stdio (stdin/stdout) using JSON-RPC protocol.
All print statements go to stderr to avoid corrupting the protocol.
"""

# ═══════════════════════════════════════════════════════════════════════════════
# IMPORTS
# ═══════════════════════════════════════════════════════════════════════════════
# Standard library imports come first, then third-party, then local modules

from __future__ import annotations  # Allows using class name in type hints before defined

import asyncio          # Async/await support for non-blocking I/O
import json             # JSON parsing for tool arguments and results
import logging          # Logging for debugging and monitoring
import os               # Environment variables (HF_MODEL)
import re               # Regular expressions for text processing
import sqlite3          # SQLite database for customers/orders/tickets
import subprocess       # Not used, but available for future extensions
import sys              # System-specific parameters (stderr, exit)
from datetime import datetime   # Timestamps for logging and tickets
from pathlib import Path        # Cross-platform file path handling
from typing import Any, Dict, List, Optional  # Type hints for better code clarity

# ─────────────────────────────────────────────────────────────────────────────
# MCP Library Imports
# FastMCP is the recommended high-level framework for MCP servers.
# It auto-generates tool schemas from type hints and docstrings.
# Install with: pip install mcp
# ─────────────────────────────────────────────────────────────────────────────
try:
    from mcp.server.fastmcp import FastMCP
    from mcp.types import TextContent
    import mcp.types as types
except ImportError:
    print("MCP not installed. Install with: pip install mcp")
    sys.exit(1)

# ─────────────────────────────────────────────────────────────────────────────
# Vector Database and PDF Processing
# ChromaDB: Vector database for semantic search (RAG retrieval)
# pypdf: PDF text extraction for loading knowledge base
# ─────────────────────────────────────────────────────────────────────────────
import chromadb                 # Vector database for embeddings and semantic search
from chromadb.config import Settings  # ChromaDB configuration
import pypdf                    # PDF text extraction

# ─────────────────────────────────────────────────────────────────────────────
# Logging Configuration
# Logs to both file and console for debugging
# Log file: mcp_server.log (in current directory)
# ─────────────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mcp_server.log'),   # Persistent log file
        logging.StreamHandler()                  # Console output (goes to stderr)
    ]
)
logger = logging.getLogger("omnitech-support-mcp")

# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ SECTION 1: CONFIGURATION AND CONSTANTS                                   ║
# ║                                                                          ║
# ║ Purpose: Define all configurable paths, settings, and mappings           ║
# ║                                                                          ║
# ║ These constants control where the server finds its data:                 ║
# ║   - Knowledge base PDFs (for RAG retrieval)                              ║
# ║   - Customer database (SQLite)                                           ║
# ║   - Initial seed data (JSON)                                             ║
# ║   - LLM model configuration                                              ║
# ║                                                                          ║
# ║ BEST PRACTICE: Configuration at the top makes it easy to modify          ║
# ║ behavior without searching through code. Consider moving to              ║
# ║ environment variables or a config file for production deployments.       ║
# ╚══════════════════════════════════════════════════════════════════════════╝

# ─────────────────────────────────────────────────────────────────────────────
# File Paths
# Using Path objects (not strings) for cross-platform compatibility
# ─────────────────────────────────────────────────────────────────────────────
KNOWLEDGE_BASE_DIR = Path("knowledge_base_pdfs")  # Directory containing PDF docs
CUSTOMER_DB_PATH = Path("customers.db")           # SQLite database file
SEED_DATA_PATH = Path("seed_data.json")           # Initial data for empty database

# ─────────────────────────────────────────────────────────────────────────────
# LLM Configuration
# This is exposed as an MCP resource (config://llm) so the RAG agent can
# discover which model to use. The model name can be overridden via
# environment variable: export HF_MODEL="your-model-name"
# ─────────────────────────────────────────────────────────────────────────────
LLM_CONFIG = {
    "model_name": os.environ.get("HF_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
    "provider": "huggingface",
    "inference_endpoint": "https://api-inference.huggingface.co"
}

# ─────────────────────────────────────────────────────────────────────────────
# Document Category Mappings
# Maps each PDF filename to its support category for classification.
# When PDFs are loaded, this mapping determines which category filter
# to apply during semantic search. This enables category-specific RAG.
# ─────────────────────────────────────────────────────────────────────────────
DOCUMENT_CATEGORIES = {
    "account_security": ["OmniTech_Account_Security_Handbook.pdf"],
    "device_troubleshooting": ["OmniTech_Device_Troubleshooting_Manual.pdf"],
    "shipping_inquiry": ["OmniTech_Global_Shipping_Logistics.pdf"],
    "returns_refunds": ["OmniTech_Returns_Policy_2024.pdf"],
}

# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ SECTION 2: CANONICAL QUERY DEFINITIONS                                   ║
# ║                                                                          ║
# ║ Purpose: Define support categories, prompt templates, and keywords       ║
# ║                                                                          ║
# ║ Each category contains:                                                  ║
# ║   - description: Human-readable description of the category              ║
# ║   - prompt_template: Template for LLM prompts (with {knowledge} and      ║
# ║     {query} placeholders that get filled in by the RAG agent)            ║
# ║   - example_queries: Sample questions for this category                  ║
# ║   - keywords: Words/phrases that indicate this category (used by         ║
# ║     the classify_query tool to match incoming queries)                   ║
# ║                                                                          ║
# ║ HOW CLASSIFICATION WORKS:                                                ║
# ║   1. User asks: "How do I reset my password?"                            ║
# ║   2. classify_query tool scans keywords for each category                ║
# ║   3. "password" and "reset" match account_security keywords              ║
# ║   4. Returns: {"suggested_query": "account_security", "confidence": 1.0} ║
# ║   5. RAG agent uses the account_security prompt_template                 ║
# ║                                                                          ║
# ║ TO ADD A NEW CATEGORY:                                                   ║
# ║   1. Add a new entry to CANONICAL_QUERIES below                          ║
# ║   2. Add corresponding PDF to DOCUMENT_CATEGORIES above                  ║
# ║   3. Place the PDF in knowledge_base_pdfs/ directory                     ║
# ╚══════════════════════════════════════════════════════════════════════════╝

CANONICAL_QUERIES = {
    "account_security": {
        "description": "Account security, passwords, 2FA, and account recovery",
        "prompt_template": """You are an OmniTech customer support specialist focused on account security.

Based on the following documentation:
{knowledge}

Please help the customer with their security-related question: {query}

Provide a clear, helpful response that:
1. Directly addresses their security concern
2. Includes step-by-step instructions if applicable
3. Emphasizes security best practices
4. Offers additional security tips when relevant

Be professional, reassuring, and thorough.""",
        "example_queries": [
            "How do I reset my password?",
            "Can you help me set up two-factor authentication?",
            "What should I do if my account is compromised?",
        ],
        "keywords": ["password", "reset", "security", "2fa", "two-factor",
                    "authentication", "account", "compromised", "hack", "secure",
                    "login", "sign in", "access", "locked out"]
    },

    "device_troubleshooting": {
        "description": "Technical issues, device problems, and troubleshooting",
        "prompt_template": """You are an OmniTech technical support specialist.

Based on the following troubleshooting documentation:
{knowledge}

Please help the customer with their device issue: {query}

Provide a clear troubleshooting response that:
1. Identifies the likely cause of the problem
2. Offers step-by-step troubleshooting instructions
3. Suggests preventive measures
4. Indicates when professional repair might be needed

Be patient, technical but accessible, and thorough.""",
        "example_queries": [
            "My device won't turn on",
            "How do I perform a factory reset?",
            "The screen is frozen",
        ],
        "keywords": ["device", "won't", "turn on", "frozen", "screen",
                    "not working", "broken", "fix", "troubleshoot", "problem",
                    "issue", "error", "crash", "slow", "battery", "overheat"]
    },

    "shipping_inquiry": {
        "description": "Order tracking, delivery times, and shipping information",
        "prompt_template": """You are an OmniTech shipping and logistics specialist.

Based on the following shipping documentation:
{knowledge}

Please help the customer with their shipping question: {query}

Provide helpful shipping information that:
1. Answers their specific shipping question
2. Provides relevant timeframes and policies
3. Explains tracking and delivery options
4. Mentions any special considerations

Be informative, precise with timeframes, and helpful.""",
        "example_queries": [
            "When will my order arrive?",
            "Do you ship internationally?",
            "How can I track my package?",
        ],
        "keywords": ["ship", "shipping", "delivery", "track", "order", "arrive",
                   "package", "international", "cost", "when will", "tracking"]
    },

    "returns_refunds": {
        "description": "Return policies, refunds, warranties, and exchanges",
        "prompt_template": """You are an OmniTech returns and warranty specialist.

Based on the following returns policy documentation:
{knowledge}

Please help the customer with their returns/warranty question: {query}

Provide a clear response about returns that:
1. Explains the relevant policy clearly
2. States timeframes and conditions
3. Describes the return/refund process
4. Mentions warranty coverage if applicable

Be understanding, clear about policies, and helpful.""",
        "example_queries": [
            "What is your return policy?",
            "How long do I have to return a product?",
            "Is my device still under warranty?",
        ],
        "keywords": ["return", "refund", "warranty", "exchange", "policy",
                  "money back", "defective", "replace", "damaged"]
    },

    "general_support": {
        "description": "General customer assistance and other inquiries",
        "prompt_template": """You are an OmniTech customer support representative.

Based on the following documentation:
{knowledge}

Please help the customer with their question: {query}

Provide a helpful, professional response that addresses their needs.""",
        "example_queries": [
            "How do I contact customer support?",
            "Tell me about OmniTech products",
        ],
        "keywords": ["help", "support", "contact", "information", "about", "general"]
    }
}

# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ SECTION 3: MCP SERVER CLASS                                              ║
# ║                                                                          ║
# ║ Purpose: Main server class that manages data and state                   ║
# ║                                                                          ║
# ║ This class encapsulates all server state and data management:            ║
# ║   - Database management (SQLite for customers, orders, tickets)          ║
# ║   - Knowledge base management (ChromaDB for vector search)               ║
# ║                                                                          ║
# ║ Tool handlers are registered as @mcp.tool() functions below the class.   ║
# ║ FastMCP auto-generates schemas from type hints and docstrings.           ║
# ║                                                                          ║
# ║ INITIALIZATION FLOW:                                                     ║
# ║   1. Create OmniTechSupport instance                                     ║
# ║   2. Set up SQLite database (creates tables, seeds data if empty)        ║
# ║   3. Set up ChromaDB and load PDF documents                              ║
# ║   4. Tools are registered via @mcp.tool() decorators (outside the class) ║
# ║   5. Resources are registered via @mcp.resource() decorators             ║
# ╚══════════════════════════════════════════════════════════════════════════╝

class OmniTechSupport:
    """
    Data and state manager for the OmniTech Customer Support MCP server.

    Manages the knowledge base (ChromaDB), customer database (SQLite),
    and request logging. Tool handlers reference this instance for data access.
    """

    def __init__(self):
        """
        Initialize the server with all required components.

        Sets up:
        - SQLite database for customer/order/ticket data
        - ChromaDB for semantic search of knowledge base
        - Request log for debugging
        """
        self.knowledge_base = None                # ChromaDB collection (set in _setup_knowledge_base)
        self.chroma_client = None                 # ChromaDB client instance
        self.request_log = []                     # Log of recent tool calls (for debugging)
        self.db_path = CUSTOMER_DB_PATH           # Path to SQLite database

        # Initialize components in order (database must exist before knowledge base)
        self._setup_database()      # 1. Create SQLite tables and seed data
        self._setup_knowledge_base()  # 2. Load PDFs into vector store

    # ─────────────────────────────────────────────────────────────────────────
    # SQLite Database Setup Methods
    #
    # These methods handle creating and managing the SQLite database that
    # stores customer information, orders, and support tickets.
    #
    # Database Schema:
    #   customers: email (PK), name, tier, support_tickets, created_at
    #   orders: id (PK), customer_email (FK), order_date, product, status
    #   tickets: id (PK), customer_email (FK), issue_type, description,
    #            priority, status, created_at, assigned_agent
    # ─────────────────────────────────────────────────────────────────────────

    def _setup_database(self):
        """
        Initialize SQLite database with customers, orders, and tickets tables.

        This method:
        1. Creates the database file if it doesn't exist
        2. Creates the three tables (customers, orders, tickets)
        3. Seeds the database with initial data from seed_data.json if empty

        The seed data pattern allows easy customization of demo data without
        modifying code. See seed_data.json and SEED_DATA_README.md for details.
        """
        logger.info(f"Initializing customer database: {self.db_path}")

        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Create customers table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS customers (
                email TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                tier TEXT DEFAULT 'Standard',
                support_tickets INTEGER DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        # Create orders table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS orders (
                id TEXT PRIMARY KEY,
                customer_email TEXT NOT NULL,
                order_date TEXT NOT NULL,
                product TEXT NOT NULL,
                status TEXT DEFAULT 'Processing',
                FOREIGN KEY (customer_email) REFERENCES customers(email)
            )
        """)

        # Create tickets table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tickets (
                id TEXT PRIMARY KEY,
                customer_email TEXT NOT NULL,
                issue_type TEXT NOT NULL,
                description TEXT,
                priority TEXT DEFAULT 'medium',
                status TEXT DEFAULT 'Open',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                assigned_agent TEXT,
                FOREIGN KEY (customer_email) REFERENCES customers(email)
            )
        """)

        # Check if we need to seed data
        cursor.execute("SELECT COUNT(*) FROM customers")
        if cursor.fetchone()[0] == 0:
            self._seed_database(cursor)

        conn.commit()
        conn.close()
        logger.info("Customer database ready")

    def _seed_database(self, cursor):
        """Seed the database from seed_data.json file."""
        logger.info("Seeding customer database...")

        if not SEED_DATA_PATH.exists():
            logger.warning(f"Seed data file not found: {SEED_DATA_PATH}")
            return

        with open(SEED_DATA_PATH, 'r') as f:
            seed_data = json.load(f)

        # Insert customers
        customers = [
            (c["email"], c["name"], c.get("tier", "Standard"), c.get("support_tickets", 0))
            for c in seed_data.get("customers", [])
        ]
        cursor.executemany(
            "INSERT INTO customers (email, name, tier, support_tickets) VALUES (?, ?, ?, ?)",
            customers
        )

        # Insert orders
        orders = [
            (o["id"], o["customer_email"], o["order_date"], o["product"], o.get("status", "Processing"))
            for o in seed_data.get("orders", [])
        ]
        cursor.executemany(
            "INSERT INTO orders (id, customer_email, order_date, product, status) VALUES (?, ?, ?, ?, ?)",
            orders
        )

        logger.info(f"Seeded {len(customers)} customers and {len(orders)} orders from {SEED_DATA_PATH}")

    def _get_db_connection(self):
        """Get a database connection."""
        return sqlite3.connect(self.db_path)

    def _get_customer_count(self) -> int:
        """Get the number of customers in the database."""
        conn = self._get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM customers")
        count = cursor.fetchone()[0]
        conn.close()
        return count

    def _get_ticket_count(self) -> int:
        """Get the total number of tickets."""
        conn = self._get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM tickets")
        count = cursor.fetchone()[0]
        conn.close()
        return count

    def _get_open_ticket_count(self) -> int:
        """Get the number of open tickets."""
        conn = self._get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM tickets WHERE status = 'Open'")
        count = cursor.fetchone()[0]
        conn.close()
        return count

    def _generate_ticket_id(self) -> str:
        """Generate a unique ticket ID."""
        count = self._get_ticket_count()
        return f"TKT-{datetime.now().strftime('%Y%m%d')}-{count + 1:04d}"

    # ─────────────────────────────────────────────────────────────────────────
    # Knowledge Base Setup Methods
    #
    # These methods handle loading PDF documents and creating the vector store
    # for semantic search (the "R" in RAG - Retrieval Augmented Generation).
    #
    # HOW RAG RETRIEVAL WORKS:
    #   1. PDFs are loaded and text is extracted using pypdf
    #   2. Text is stored in ChromaDB with category metadata
    #   3. ChromaDB automatically creates embeddings for each document
    #   4. When a query comes in, ChromaDB finds similar documents using
    #      cosine similarity between query embedding and document embeddings
    #   5. Retrieved documents are used as context for the LLM response
    #
    # ChromaDB uses a default embedding model (all-MiniLM-L6-v2) which is
    # downloaded automatically on first run. For production, consider using
    # a more powerful embedding model.
    # ─────────────────────────────────────────────────────────────────────────

    def _load_pdf_documents(self) -> List[Dict]:
        """Load and parse PDF documents from knowledge base directory."""
        documents = []

        if not KNOWLEDGE_BASE_DIR.exists():
            logger.error(f"Knowledge base directory not found: {KNOWLEDGE_BASE_DIR}")
            return documents

        for filename in os.listdir(KNOWLEDGE_BASE_DIR):
            if not filename.endswith('.pdf'):
                continue

            file_path = KNOWLEDGE_BASE_DIR / filename
            filename_without_ext = filename.replace('.pdf', '')

            # Determine category from filename
            category = "general_support"
            for cat, files in DOCUMENT_CATEGORIES.items():
                if filename in files:
                    category = cat
                    break

            try:
                with open(file_path, 'rb') as f:
                    pdf_reader = pypdf.PdfReader(f)
                    page_count = 0
                    for page_num, page in enumerate(pdf_reader.pages, start=1):
                        page_text = page.extract_text() or ""
                        page_text = re.sub(r'\s+', ' ', page_text.strip())

                        # Skip pages with very little text
                        if len(page_text) < 50:
                            continue

                        documents.append({
                            "id": f"{filename_without_ext}_page{page_num}",
                            "text": page_text,
                            "category": category,
                            "source": filename,
                            "page": page_num
                        })
                        page_count += 1

                logger.info(f"Loaded: {filename} -> {category} ({page_count} pages)")

            except Exception as e:
                logger.error(f"Failed to load {filename}: {e}")

        return documents

    def _setup_knowledge_base(self):
        """Initialize ChromaDB and load documents."""
        logger.info("Initializing knowledge base...")

        self.chroma_client = chromadb.Client(Settings(anonymized_telemetry=False))

        try:
            self.chroma_client.delete_collection("omnitech_knowledge")
        except:
            pass

        self.knowledge_base = self.chroma_client.create_collection("omnitech_knowledge")

        documents = self._load_pdf_documents()

        for doc in documents:
            self.knowledge_base.add(
                documents=[doc["text"]],
                metadatas=[{
                    "category": doc["category"],
                    "source": doc["source"],
                    "page": doc["page"]
                }],
                ids=[doc["id"]]
            )

        logger.info(f"Knowledge base ready: {self.knowledge_base.count()} chunks")

    # ─────────────────────────────────────────────────────────────────────────
    # Request Logging
    #
    # Logs all tool calls for debugging and analytics. The request log is
    # returned by get_server_stats and displayed in the MCP Monitor tab.
    # Keeps only the last 50 entries to avoid memory issues.
    # ─────────────────────────────────────────────────────────────────────────

    def log_request(self, tool_name: str, args: Dict[str, Any], result: Any):
        """Log MCP tool requests."""
        log_entry = {
            "timestamp": datetime.now().isoformat(),
            "tool": tool_name,
            "arguments": args,
            "result_preview": str(result)[:200]
        }
        self.request_log.append(log_entry)

        if len(self.request_log) > 50:
            self.request_log = self.request_log[-50:]

        logger.info(f"Tool call: {tool_name}")


# ═══════════════════════════════════════════════════════════════════════════════
# Create FastMCP server and data manager instances
# ═══════════════════════════════════════════════════════════════════════════════

mcp = FastMCP("omnitech-support")
server_data = OmniTechSupport()


# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ TOOL DEFINITIONS (using FastMCP @mcp.tool() decorators)                  ║
# ║                                                                          ║
# ║ FastMCP auto-generates tool schemas from:                                ║
# ║   - Function name → tool name                                           ║
# ║   - Docstring → tool description (shown to the LLM for tool selection)  ║
# ║   - Type hints → input schema (JSON Schema)                             ║
# ║   - Default values → optional parameters                                ║
# ║                                                                          ║
# ║ The LLM uses the tool descriptions and schemas to decide which tool      ║
# ║ to call and with what arguments. Good docstrings are essential!          ║
# ╚══════════════════════════════════════════════════════════════════════════╝


# ─────────────────────────────────────────────────────────────────────────────
# Classification Tools
#
# These tools implement query classification — the first step in the RAG
# pipeline. They determine what type of support request the user is making
# so we can retrieve the right documents.
# ─────────────────────────────────────────────────────────────────────────────

@mcp.tool()
def classify_query(user_query: str) -> str:
    """Classify a customer query into a support category.

    Analyzes the query text against keyword patterns to determine
    the most relevant support category. Returns the category name,
    confidence score, and alternative matches.

    Args:
        user_query: The customer's query to classify
    """
    user_lower = user_query.lower()
    scores = {}

    # Score based on keyword matching (primary signal)
    for query_name, config in CANONICAL_QUERIES.items():
        score = 0
        keywords = config.get("keywords", [])

        for keyword in keywords:
            if keyword in user_lower:
                # Give higher weight to longer/more specific keywords
                score += len(keyword.split())

        scores[query_name] = score

    # Also check example queries (secondary signal, weighted lower)
    for query_name, config in CANONICAL_QUERIES.items():
        for example in config["example_queries"]:
            # Only count content words, not common words like "how", "do", "i", "my"
            stop_words = {"how", "do", "i", "my", "the", "a", "an", "is", "it", "to", "can", "you", "what"}
            example_words = set(example.lower().split()) - stop_words
            user_words = set(user_lower.split()) - stop_words
            overlap = len(example_words.intersection(user_words))
            if overlap > 0 and example_words:
                similarity = overlap / len(example_words)
                # Add to score instead of replacing (weighted at 0.5)
                scores[query_name] = scores.get(query_name, 0) + (similarity * 0.5)

    # Find best match
    if not scores or max(scores.values()) == 0:
        result = {
            "suggested_query": "general_support",
            "confidence": 0.3,
            "alternatives": [],
            "reason": "No specific category matched"
        }
    else:
        best_query = max(scores, key=scores.get)
        best_score = scores[best_query]
        # Normalize confidence: score of 2+ keywords = high confidence
        confidence = min(best_score / 2.0, 1.0)

        alternatives = [
            {"query": name, "score": round(score, 3)}
            for name, score in sorted(scores.items(), key=lambda x: x[1], reverse=True)
            if score > 0 and name != best_query
        ][:2]

        result = {
            "suggested_query": best_query,
            "confidence": round(confidence, 3),
            "alternatives": alternatives,
            "reason": f"Matched to {best_query}"
        }

    server_data.log_request("classify_query", {"user_query": user_query}, result)
    return json.dumps(result)


@mcp.tool()
def get_query_template(query_name: str) -> str:
    """Get the prompt template for a support category.

    Returns the LLM prompt template and description for the specified
    category. Templates contain {knowledge} and {query} placeholders.

    Args:
        query_name: Category name (e.g., 'account_security', 'device_troubleshooting')
    """
    if query_name not in CANONICAL_QUERIES:
        result = {"error": f"Unknown category: {query_name}"}
    else:
        config = CANONICAL_QUERIES[query_name]
        result = {
            "template": config["prompt_template"],
            "description": config["description"]
        }

    server_data.log_request("get_query_template", {"query_name": query_name}, result)
    return json.dumps(result)


@mcp.tool()
def list_categories() -> str:
    """List all available support categories.

    Returns all support category names, descriptions, and example queries.
    Useful for understanding what types of queries the system can handle.
    """
    categories = []
    for name, config in CANONICAL_QUERIES.items():
        categories.append({
            "name": name,
            "description": config["description"],
            "example_queries": config["example_queries"][:3]
        })

    result = {"categories": categories}
    server_data.log_request("list_categories", {}, result)
    return json.dumps(result)


# ─────────────────────────────────────────────────────────────────────────────
# Knowledge Tools
#
# These tools implement RAG retrieval — searching the ChromaDB vector store
# to find relevant documents for a query.
#
#   - search_knowledge: Returns detailed results with similarity scores
#     (useful for debugging and the Knowledge Search tab)
#   - get_knowledge_for_query: Returns concatenated text for LLM context
#     (optimized for use in the RAG pipeline)
# ─────────────────────────────────────────────────────────────────────────────

@mcp.tool()
def search_knowledge(query: str, category: str = "", max_results: int = 3) -> str:
    """Search the knowledge base for relevant documents.

    Performs semantic search over PDF documentation using ChromaDB.
    Returns matching documents with similarity scores and source attribution.

    Args:
        query: Search query text
        category: Optional category filter (e.g., 'account_security')
        max_results: Maximum number of results to return (default: 3)
    """
    where_clause = {"category": category} if category else None

    results = server_data.knowledge_base.query(
        query_texts=[query],
        n_results=max_results,
        where=where_clause,
        include=["documents", "metadatas", "distances"]
    )

    matches = []
    if results["documents"] and results["documents"][0]:
        for doc, meta, dist in zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0]
        ):
            matches.append({
                "content": doc[:500] + "..." if len(doc) > 500 else doc,
                "category": meta.get("category", "unknown"),
                "source": meta.get("source", "Unknown"),
                "page": meta.get("page", "N/A"),
                "distance": dist,
                "similarity": round(1 / (1 + dist), 3)
            })

    result = {
        "matches": matches,
        "count": len(matches),
        "query": query,
        "category_filter": category
    }

    server_data.log_request("search_knowledge", {"query": query, "category": category}, result)
    return json.dumps(result)


@mcp.tool()
def get_knowledge_for_query(category: str, query: str, max_results: int = 3) -> str:
    """Get concatenated knowledge text for a category and query.

    Optimized for use as LLM context in the RAG pipeline. Returns
    the full text of matching documents joined together, plus source list.

    Args:
        category: Support category to filter by
        query: The user's question
        max_results: Maximum number of documents to retrieve (default: 3)
    """
    where_clause = {"category": category} if category else None

    results = server_data.knowledge_base.query(
        query_texts=[query],
        n_results=max_results,
        where=where_clause,
        include=["documents", "metadatas"]
    )

    knowledge_parts = []
    sources = set()

    if results["documents"] and results["documents"][0]:
        for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
            knowledge_parts.append(doc)
            source = meta.get("source", "Unknown")
            page = meta.get("page", "N/A")
            sources.add((source, str(page)))

    result = {
        "knowledge": "\n\n---\n\n".join(knowledge_parts) if knowledge_parts else "No relevant documentation found.",
        "sources": [f"{s} (page {p})" for s, p in sources],
        "chunks_retrieved": len(knowledge_parts)
    }

    server_data.log_request("get_knowledge_for_query", {"category": category, "query": query}, result)
    return json.dumps(result)


# ─────────────────────────────────────────────────────────────────────────────
# Customer Tools
#
# These tools manage customer data in the SQLite database — looking up
# customer information and creating support tickets.
#
# The lookup_customer tool is essential for personalized responses. When
# a customer asks about their order, we can look up their actual orders
# and provide specific information rather than generic responses.
# ─────────────────────────────────────────────────────────────────────────────

@mcp.tool()
def lookup_customer(email: str) -> str:
    """Look up customer information by email address.

    Returns customer details (name, tier, ticket count) and their
    order history from the SQLite database.

    Args:
        email: Customer email address to look up
    """
    email_lower = email.lower()

    conn = server_data._get_db_connection()
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()

    # Get customer info
    cursor.execute(
        "SELECT email, name, tier, support_tickets FROM customers WHERE LOWER(email) = ?",
        (email_lower,)
    )
    customer_row = cursor.fetchone()

    if customer_row:
        # Get customer's orders
        cursor.execute(
            "SELECT id, order_date, product, status FROM orders WHERE LOWER(customer_email) = ?",
            (email_lower,)
        )
        orders = [
            {"id": row["id"], "date": row["order_date"], "product": row["product"], "status": row["status"]}
            for row in cursor.fetchall()
        ]

        result = {
            "found": True,
            "email": customer_row["email"],
            "name": customer_row["name"],
            "tier": customer_row["tier"],
            "support_tickets": customer_row["support_tickets"],
            "orders": orders
        }
    else:
        result = {
            "found": False,
            "email": email,
            "message": "Customer not found in database"
        }

    conn.close()
    server_data.log_request("lookup_customer", {"email": email}, result)
    return json.dumps(result)


@mcp.tool()
def create_support_ticket(customer_email: str, issue_type: str, description: str, priority: str = "medium") -> str:
    """Create a new support ticket in the database.

    Creates a ticket with auto-generated ID and timestamp.
    Updates the customer's ticket count.

    Args:
        customer_email: Customer's email address
        issue_type: Type of issue (e.g., 'account_security', 'device_troubleshooting')
        description: Description of the issue
        priority: Priority level ('low', 'medium', 'high') — default: 'medium'
    """
    ticket_id = server_data._generate_ticket_id()
    created_at = datetime.now().isoformat()

    conn = server_data._get_db_connection()
    cursor = conn.cursor()

    # Insert ticket into database
    cursor.execute("""
        INSERT INTO tickets (id, customer_email, issue_type, description, priority, status, created_at)
        VALUES (?, ?, ?, ?, ?, 'Open', ?)
    """, (ticket_id, customer_email, issue_type, description, priority, created_at))

    # Update customer's ticket count if they exist
    cursor.execute(
        "UPDATE customers SET support_tickets = support_tickets + 1 WHERE LOWER(email) = ?",
        (customer_email.lower(),)
    )

    conn.commit()
    conn.close()

    ticket = {
        "id": ticket_id,
        "customer_email": customer_email,
        "issue_type": issue_type,
        "description": description,
        "priority": priority,
        "status": "Open",
        "created_at": created_at,
        "assigned_agent": None
    }

    logger.info(f"Created ticket {ticket_id} for {customer_email}")
    server_data.log_request("create_support_ticket", {"customer_email": customer_email, "issue_type": issue_type}, ticket)
    return json.dumps(ticket)


# ─────────────────────────────────────────────────────────────────────────────
# Statistics and Ticket Tools
#
# These tools provide server statistics and ticket management.
# The get_server_stats tool is useful for the MCP Monitor tab in the UI
# and for debugging.
# ─────────────────────────────────────────────────────────────────────────────

@mcp.tool()
def get_server_stats() -> str:
    """Get MCP server statistics and recent activity.

    Returns server health info, document counts, customer counts,
    available tools, and recent request history.
    """
    stats = {
        "server_status": "Active",
        "llm_model": LLM_CONFIG["model_name"],
        "knowledge_chunks": server_data.knowledge_base.count() if server_data.knowledge_base else 0,
        "customers_in_db": server_data._get_customer_count(),
        "database_type": "SQLite",
        "database_path": str(server_data.db_path),
        "total_tickets": server_data._get_ticket_count(),
        "open_tickets": server_data._get_open_ticket_count(),
        "support_categories": list(CANONICAL_QUERIES.keys()),
        "total_requests": len(server_data.request_log),
        "recent_requests": server_data.request_log[-10:],
        "tools_available": [
            "classify_query",
            "get_query_template",
            "list_categories",
            "search_knowledge",
            "get_knowledge_for_query",
            "lookup_customer",
            "create_support_ticket",
            "get_tickets",
            "get_server_stats"
        ]
    }

    server_data.log_request("get_server_stats", {}, stats)
    return json.dumps(stats, indent=2)


@mcp.tool()
def get_tickets(customer_email: str = "", status: str = "", limit: int = 10) -> str:
    """Get support tickets with optional filters.

    Retrieves tickets from the database, optionally filtered by
    customer email and/or status.

    Args:
        customer_email: Filter by customer email (optional)
        status: Filter by status — 'Open', 'Closed', etc. (optional)
        limit: Maximum number of tickets to return (default: 10)
    """
    conn = server_data._get_db_connection()
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()

    # Build query with optional filters
    query = "SELECT * FROM tickets WHERE 1=1"
    params = []

    if customer_email:
        query += " AND LOWER(customer_email) = ?"
        params.append(customer_email.lower())

    if status:
        query += " AND status = ?"
        params.append(status)

    query += " ORDER BY created_at DESC LIMIT ?"
    params.append(limit)

    cursor.execute(query, params)
    rows = cursor.fetchall()

    tickets = [
        {
            "id": row["id"],
            "customer_email": row["customer_email"],
            "issue_type": row["issue_type"],
            "description": row["description"],
            "priority": row["priority"],
            "status": row["status"],
            "created_at": row["created_at"],
            "assigned_agent": row["assigned_agent"]
        }
        for row in rows
    ]

    conn.close()

    result = {
        "tickets": tickets,
        "count": len(tickets),
        "filters": {
            "customer_email": customer_email,
            "status": status,
            "limit": limit
        }
    }

    server_data.log_request("get_tickets", {"customer_email": customer_email, "status": status}, result)
    return json.dumps(result, indent=2)


# ─────────────────────────────────────────────────────────────────────────────
# Resource Registration
#
# MCP Resources are read-only data that clients can access. Unlike tools
# (which perform actions), resources just expose data. Resources are
# identified by URIs like "config://llm" or "data://tickets".
#
# RESOURCES IN THIS SERVER:
#   config://llm        - LLM model configuration
#   config://database   - Database statistics
#   config://categories - Support category definitions
#   data://tickets      - Current support tickets
#
# The RAG agent uses the config://llm resource to discover which
# LLM model to use for generating responses.
# ─────────────────────────────────────────────────────────────────────────────

@mcp.resource("config://llm")
def get_llm_config() -> str:
    """Current LLM model configuration."""
    return json.dumps({
        "model_name": LLM_CONFIG["model_name"],
        "provider": LLM_CONFIG["provider"],
        "inference_endpoint": LLM_CONFIG["inference_endpoint"],
        "description": "HuggingFace Inference API with Llama model"
    }, indent=2)


@mcp.resource("config://database")
def get_database_config() -> str:
    """Customer database configuration and statistics."""
    conn = server_data._get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) FROM customers")
    customer_count = cursor.fetchone()[0]
    cursor.execute("SELECT COUNT(*) FROM orders")
    order_count = cursor.fetchone()[0]
    cursor.execute("SELECT COUNT(*) FROM tickets")
    ticket_count = cursor.fetchone()[0]
    conn.close()

    return json.dumps({
        "type": "SQLite",
        "path": str(server_data.db_path),
        "tables": ["customers", "orders", "tickets"],
        "customer_count": customer_count,
        "order_count": order_count,
        "ticket_count": ticket_count,
        "description": "SQLite database for customer, order, and ticket information"
    }, indent=2)


@mcp.resource("config://categories")
def get_categories_config() -> str:
    """Available support categories and their configurations."""
    categories = {}
    for name, config in CANONICAL_QUERIES.items():
        categories[name] = {
            "description": config["description"],
            "example_queries": config["example_queries"],
            "keyword_count": len(config.get("keywords", []))
        }
    return json.dumps({
        "categories": categories,
        "total_categories": len(categories)
    }, indent=2)


@mcp.resource("data://tickets")
def get_tickets_data() -> str:
    """Current support tickets in the system."""
    conn = server_data._get_db_connection()
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    cursor.execute("""
        SELECT id, customer_email, issue_type, description, priority, status, created_at, assigned_agent
        FROM tickets ORDER BY created_at DESC LIMIT 50
    """)
    rows = cursor.fetchall()
    conn.close()

    tickets = [
        {
            "id": row["id"],
            "customer_email": row["customer_email"],
            "issue_type": row["issue_type"],
            "description": row["description"],
            "priority": row["priority"],
            "status": row["status"],
            "created_at": row["created_at"],
            "assigned_agent": row["assigned_agent"]
        }
        for row in rows
    ]

    return json.dumps({
        "tickets": tickets,
        "total_count": len(tickets),
        "open_count": len([t for t in tickets if t["status"] == "Open"])
    }, indent=2)


# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ SECTION 4: MAIN ENTRY POINT                                              ║
# ║                                                                          ║
# ║ Purpose: Start the MCP server and handle the stdio communication         ║
# ║                                                                          ║
# ║ FastMCP handles all the protocol details:                                ║
# ║   - stdio transport (JSON-RPC over stdin/stdout)                         ║
# ║   - Tool schema generation from type hints                               ║
# ║   - Tool dispatch to the right function                                  ║
# ║   - Resource URI routing                                                 ║
# ║                                                                          ║
# ║ IMPORTANT: Never print to stdout! It will corrupt the JSON-RPC protocol. ║
# ║ Always use: print("message", file=sys.stderr) or logger.info()           ║
# ╚══════════════════════════════════════════════════════════════════════════╝

if __name__ == "__main__":
    # ─────────────────────────────────────────────────────────────────────────
    # Startup banner - goes to stderr because stdout is reserved for MCP
    # ─────────────────────────────────────────────────────────────────────────
    print("=" * 60, file=sys.stderr)
    print("OmniTech Customer Support MCP Server", file=sys.stderr)
    print("=" * 60, file=sys.stderr)
    print(f"Knowledge base: {KNOWLEDGE_BASE_DIR}", file=sys.stderr)
    print(f"Customer database: {CUSTOMER_DB_PATH}", file=sys.stderr)
    print(f"LLM model: {LLM_CONFIG['model_name']}", file=sys.stderr)
    print(f"Support categories: {list(CANONICAL_QUERIES.keys())}", file=sys.stderr)
    print("=" * 60, file=sys.stderr)

    # FastMCP handles stdio transport, tool dispatch, and resource routing
    mcp.run()