"""DataBridge AI - MCP Server Implementation.

A headless, MCP-native data reconciliation engine that bridges messy sources
(OCR/PDF/SQL) with a structured Python-based comparison pipeline.
"""
from fastmcp import FastMCP
import pandas as pd
import hashlib
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional
import re

# Conditional imports for optional dependencies
try:
    from rapidfuzz import fuzz, process
    RAPIDFUZZ_AVAILABLE = True
except ImportError:
    RAPIDFUZZ_AVAILABLE = False

try:
    from sqlalchemy import create_engine, text
    SQLALCHEMY_AVAILABLE = True
except ImportError:
    SQLALCHEMY_AVAILABLE = False

try:
    import pytesseract
    from PIL import Image
    TESSERACT_AVAILABLE = True
except ImportError:
    TESSERACT_AVAILABLE = False

try:
    from pypdf import PdfReader
    PYPDF_AVAILABLE = True
except ImportError:
    PYPDF_AVAILABLE = False

# Turbo engine (Polars + DuckDB acceleration)
try:
    from src.turbo import (
        read_csv as turbo_read_csv,
        dataframe_from_records as turbo_from_records,
        profile_stats as turbo_profile_stats,
        compute_row_hashes_vectorized,
        query_local_sql,
        register_table as turbo_register_table,
        list_tables as turbo_list_tables,
        export_to_parquet as turbo_export_to_parquet,
        engine_status as turbo_engine_status,
        POLARS_AVAILABLE as TURBO_POLARS,
        DUCKDB_AVAILABLE as TURBO_DUCKDB,
    )
    TURBO_AVAILABLE = True
except ImportError:
    try:
        from turbo import (
            read_csv as turbo_read_csv,
            dataframe_from_records as turbo_from_records,
            profile_stats as turbo_profile_stats,
            compute_row_hashes_vectorized,
            query_local_sql,
            register_table as turbo_register_table,
            list_tables as turbo_list_tables,
            export_to_parquet as turbo_export_to_parquet,
            engine_status as turbo_engine_status,
            POLARS_AVAILABLE as TURBO_POLARS,
            DUCKDB_AVAILABLE as TURBO_DUCKDB,
        )
        TURBO_AVAILABLE = True
    except ImportError:
        TURBO_AVAILABLE = False
        TURBO_POLARS = False
        TURBO_DUCKDB = False

# Handle imports for both module and direct execution
try:
    from src.config import settings
except ImportError:
    from config import settings

# License Management - Import the plugin system
try:
    from src.plugins import get_license_manager, LicenseManager
    LICENSE_MANAGER = get_license_manager()
except ImportError:
    try:
        from plugins import get_license_manager, LicenseManager
        LICENSE_MANAGER = get_license_manager()
    except ImportError:
        # Fallback: No license management - run as full version (for backward compatibility)
        LICENSE_MANAGER = None
        print("[License] Plugin system not available - running in full mode")

# Determine edition based on license
if LICENSE_MANAGER:
    _edition = LICENSE_MANAGER.tier
    _edition_msg = LICENSE_MANAGER.validation_message
    print(f"[License] {_edition_msg}")
else:
    _edition = "FULL"
    _edition_msg = "Full version (no license restrictions)"

__version__ = "0.42"
__edition__ = _edition

# Report turbo engine status at startup
if TURBO_AVAILABLE:
    _polars_mark = "YES" if TURBO_POLARS else "NO"
    _duckdb_mark = "YES" if TURBO_DUCKDB else "NO"
    print(f"[Turbo] Polars {_polars_mark} DuckDB {_duckdb_mark}")
else:
    print("[Turbo] Not available (install polars + duckdb for acceleration)")


def is_pro_feature_enabled() -> bool:
    """Check if Pro features should be loaded."""
    if LICENSE_MANAGER is None:
        return True  # No license manager = full version
    return LICENSE_MANAGER.is_pro()


def is_enterprise_feature_enabled() -> bool:
    """Check if Enterprise features should be loaded."""
    if LICENSE_MANAGER is None:
        return True  # No license manager = full version
    return LICENSE_MANAGER.is_enterprise()


# Initialize the Server
mcp = FastMCP(
    "DNX Hierarchy Manager",
    instructions="""Professional Data Reconciliation & Hierarchy Management Agent.

I help you with two major capabilities:

**Data Reconciliation:**
- Compare and validate data from CSV, SQL, PDF, and JSON sources
- Find orphans, conflicts, and fuzzy matches between datasets
- Profile data quality and detect schema drift

**Hierarchy Builder:**
- Create and manage multi-level hierarchy projects (up to 15 levels)
- Define source mappings linking database columns to hierarchy nodes
- Build calculation formulas (SUM, SUBTRACT, MULTIPLY, DIVIDE)
- Export hierarchies to CSV/JSON and generate deployment scripts
- Deploy hierarchies to Snowflake and other databases
- Model application configuration as hierarchy-governed datasets
- Publish/pull hierarchy-managed config to/from Snowflake APP_CONFIG with audit

Use me for data quality checks, building financial hierarchies, and managing complex data structures."""
)

# Ensure data directory exists
Path(settings.data_dir).mkdir(parents=True, exist_ok=True)

# Initialize audit log with headers if it doesn't exist
if not Path(settings.audit_log).exists():
    with open(settings.audit_log, "w") as f:
        f.write("timestamp,user,action,impact\n")

# Initialize workflow file if it doesn't exist
if not Path(settings.workflow_file).exists():
    with open(settings.workflow_file, "w") as f:
        json.dump({"version": "1.0", "steps": []}, f)


# =============================================================================
# Internal Helpers
# =============================================================================

def log_action(user: str, action: str, impact: str) -> None:
    """Record an action to the audit trail (no PII)."""
    timestamp = datetime.now().isoformat()
    # Sanitize to prevent CSV injection
    action = action.replace(",", ";").replace("\n", " ")[:100]
    impact = impact.replace(",", ";").replace("\n", " ")[:200]
    log_entry = f"{timestamp},{user},{action},{impact}\n"
    with open(settings.audit_log, "a") as f:
        f.write(log_entry)


def _load_trust_public_key_registry() -> Dict[str, str]:
    path = getattr(settings, "trust_public_key_registry_path", Path("data/datashield/trust_public_keys.json"))
    p = Path(path)
    if not p.exists():
        return {}
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
        if isinstance(data, dict):
            return {str(k): str(v) for k, v in data.items()}
    except Exception:
        return {}
    return {}


def _enforce_trust_for_ai_tool(
    *,
    tool_name: str,
    lane: str,
    trust_attestation: str = "",
) -> Optional[Dict[str, Any]]:
    """Enforce trust attestation for protected AI tool routes.

    Returns an error dict when blocked, otherwise None.
    """
    if not getattr(settings, "trust_enforcement_enabled", True):
        return None

    try:
        from src.datashield.policy_engine import enforce_attestation
        from src.datashield.types import TrustLane
        from src.datashield.attestation import persist_attestation_event
    except Exception as exc:
        return {"error": f"Trust enforcement unavailable: {exc}"}

    att_obj = None
    if trust_attestation:
        try:
            att_obj = json.loads(trust_attestation)
        except json.JSONDecodeError:
            return {"error": "Invalid trust_attestation JSON", "code": "TRUST_ATTESTATION_INVALID_FORMAT"}

    registry = _load_trust_public_key_registry()
    result = enforce_attestation(
        att_obj,
        lane=TrustLane(lane),
        public_key_registry=registry,
    )
    event = {
        "tool": tool_name,
        "lane": lane,
        "result": result,
        "timestamp": datetime.now().isoformat(),
    }
    try:
        persist_attestation_event(event, base_dir=str(Path(settings.data_dir) / "datashield"))
    except Exception:
        pass

    if result.get("ok"):
        return None

    if getattr(settings, "trust_enforcement_mode", "hard_fail") == "warn":
        log_action("TRUST", "trust_warn", f"{tool_name}:{result.get('code', 'unknown')}")
        return None
    return {"error": "Trust enforcement failed", "code": result.get("code", "TRUST_ENFORCEMENT_FAILED")}


def _coerce_trust_attestation_str(
    explicit_attestation: str = "",
    arguments_dict: Optional[Dict[str, Any]] = None,
) -> str:
    """Resolve trust attestation from explicit arg or tool payload aliases.

    Priority:
    1) explicit_attestation argument if provided
    2) arguments_dict["trust_attestation"] if provided
    3) arguments_dict["_trust_attestation"] dict serialized as JSON
    """
    if explicit_attestation:
        return explicit_attestation

    arguments_dict = arguments_dict or {}
    payload_att = arguments_dict.get("trust_attestation")
    if isinstance(payload_att, str) and payload_att.strip():
        return payload_att
    if isinstance(payload_att, dict):
        return json.dumps(payload_att)

    alias_att = arguments_dict.get("_trust_attestation")
    if isinstance(alias_att, dict):
        return json.dumps(alias_att)
    if isinstance(alias_att, str) and alias_att.strip():
        return alias_att
    return ""


def _resolve_trust_lane_for_tool(tool_name: str) -> Optional[str]:
    """Return required trust lane for a tool routed via gateway run_tool().

    Returns:
    - "databridge_ai_masked" for DataBridge AI routes.
    - "client_cortex_raw" for client-owned Cortex routes.
    - None when no trust gate is required.
    """
    name = (tool_name or "").strip().lower()
    if not name:
        return None

    # Prefix-based fast path
    if name.startswith(("cortex_", "analyst_")):
        return "client_cortex_raw"
    if name.startswith(("ai_", "plan_", "suggest_")):
        return "databridge_ai_masked"

    # Metadata-based fallback from dynamic registry
    try:
        meta = _dynamic_registry.get_tool_metadata(tool_name)
    except Exception:
        meta = None

    domain = (getattr(meta, "domain", "") or "").strip().lower()
    if domain == "cortex":
        return "client_cortex_raw"
    if domain in {"ai_discovery", "planner", "orchestrator", "recommendations"}:
        return "databridge_ai_masked"
    return None


def _annotate_tool_list_with_trust_lane(rows: Any) -> Any:
    """Attach trust_lane to tool rows when they expose a tool name."""
    if not isinstance(rows, list):
        return rows
    enriched = []
    for item in rows:
        if isinstance(item, dict) and "name" in item:
            row = dict(item)
            trust_lane = _resolve_trust_lane_for_tool(str(item.get("name", "")))
            row["trust_lane"] = trust_lane
            row["required_trust_attestation"] = trust_lane is not None
            enriched.append(row)
        else:
            enriched.append(item)
    return enriched


def _build_trust_preflight_for_tool(tool_name: str) -> Dict[str, Any]:
    """Build structured trust preflight guidance for a tool."""
    trust_lane = _resolve_trust_lane_for_tool(tool_name)
    required = trust_lane is not None
    lane_rules = {
        "databridge_ai_masked": "valid signature required; attestation.masked must be true",
        "client_cortex_raw": "valid signature required; masked may be false",
    }
    return {
        "required_trust_attestation": required,
        "trust_lane": trust_lane,
        "enforcement_mode": getattr(settings, "trust_enforcement_mode", "hard_fail"),
        "attestation_format": "json_string",
        "required_fields": [
            "dataset_id",
            "lane",
            "masked",
            "policy_version",
            "input_hash",
            "created_at",
            "public_key_id",
            "signature",
        ] if required else [],
        "lane_rule": lane_rules.get(trust_lane, "none"),
    }


def _annotate_tool_list_with_trust_preflight(rows: Any) -> Any:
    """Attach trust_preflight guidance to tool rows when they expose a tool name."""
    if not isinstance(rows, list):
        return rows
    enriched = []
    for item in rows:
        if isinstance(item, dict) and "name" in item:
            row = dict(item)
            row["trust_preflight"] = _build_trust_preflight_for_tool(str(item.get("name", "")))
            enriched.append(row)
        else:
            enriched.append(item)
    return enriched


def compute_row_hash(row: pd.Series, columns: list) -> str:
    """Compute a deterministic SHA-256 hash for a row (truncated to 16 chars)."""
    values = "|".join(str(row[col]) for col in columns)
    return hashlib.sha256(values.encode()).hexdigest()[:16]


def truncate_dataframe(df: pd.DataFrame, max_rows: int = None) -> pd.DataFrame:
    """Truncate DataFrame to respect context sensitivity rules."""
    max_rows = max_rows or settings.max_rows_display
    return df.head(max_rows)


# =============================================================================
# Phase 0: File Discovery & Staging Tools
# =============================================================================

def get_common_search_paths() -> list:
    """Get common directories where files might be located."""
    home = Path.home()
    cwd = Path.cwd()

    paths = [
        cwd,
        cwd / "data",
        cwd / "result_export",
        cwd / "uploads",
        home,
        home / "Downloads",
        home / "Documents",
        home / "Desktop",
        home / "uploads",
        Path("/tmp"),
        Path("/tmp/uploads"),
    ]

    # Add Claude-specific paths
    claude_paths = [
        home / ".claude" / "uploads",
        home / "AppData" / "Local" / "Claude" / "uploads",
        home / "AppData" / "Roaming" / "Claude" / "uploads",
        Path("C:/Users") / home.name / "AppData" / "Local" / "Temp",
    ]
    paths.extend(claude_paths)

    return [p for p in paths if p.exists()]


@mcp.tool()
def find_files(
    pattern: str = "*.csv",
    search_name: str = "",
    max_results: int = 20
) -> str:
    """
    Search for files across common directories.

    Use this tool when you can't find a file or need to discover available files.
    It searches Downloads, Documents, Desktop, temp folders, and the DataBridge
    data directory.

    Args:
        pattern: Glob pattern to match (default "*.csv"). Examples:
                 - "*.csv" for all CSV files
                 - "*.xlsx" for Excel files
                 - "*" for all files
        search_name: Optional filename substring to filter results (case-insensitive)
        max_results: Maximum number of results to return (default 20)

    Returns:
        JSON with found files, their paths, sizes, and modification times.

    Example:
        find_files(pattern="*.csv", search_name="hierarchy")
    """
    try:
        from databridge.connectors import find_files as _find_files
        result = _find_files(
            pattern=pattern, search_name=search_name,
            max_results=max_results, data_dir=settings.data_dir,
        )
        result["tip"] = "Use stage_file(source_path) to copy a file to the DataBridge data directory for easier access."
        log_action("AI_AGENT", "find_files", f"Found {result['files_found']} files matching {pattern}")
        return json.dumps(result, indent=2)

    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def stage_file(
    source_path: str,
    new_name: str = ""
) -> str:
    """
    Copy a file to the DataBridge data directory for easy access.

    Use this when you find a file with find_files() but it's in an inconvenient
    location. This copies it to the DataBridge data directory where all tools
    can easily access it.

    Args:
        source_path: Full path to the source file
        new_name: Optional new filename (keeps original name if not provided)

    Returns:
        JSON with the new file path and confirmation.

    Example:
        stage_file("/Users/john/Downloads/my_data.csv")
        stage_file("/tmp/upload123.csv", new_name="quarterly_report.csv")
    """
    try:
        from databridge.connectors import stage_file as _stage_file
        result = _stage_file(source_path=source_path, data_dir=settings.data_dir, new_name=new_name)
        result["tip"] = f"You can now use this path: {result['destination']}"
        log_action("AI_AGENT", "stage_file", f"Staged {Path(source_path).name} to {result['destination']}")
        return json.dumps(result, indent=2)

    except FileNotFoundError:
        return json.dumps({
            "error": f"File not found: {source_path}",
            "suggestion": "Use find_files() to locate the file first.",
            "working_directory": str(Path.cwd()),
        })
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def get_working_directory() -> str:
    """
    Get the current working directory and DataBridge data directory paths.

    Use this to understand where DataBridge is looking for files and where
    to place files for easy access.

    Returns:
        JSON with working directory, data directory, and available files.
    """
    try:
        from databridge.connectors import get_working_directory as _get_wd
        result = _get_wd(data_dir=settings.data_dir)
        result["tip"] = "Use find_files() to search for files, or stage_file() to copy files here."
        return json.dumps(result, indent=2)

    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 1: Data Loading Tools
# =============================================================================

@mcp.tool()
def load_csv(file_path: str, preview_rows: int = 5) -> str:
    """
    Load a CSV file and return a preview with schema information.

    Args:
        file_path: Path to the CSV file.
        preview_rows: Number of rows to preview (max 10).

    Returns:
        JSON with schema info and sample data.
    """
    try:
        from databridge.ingestion import load_csv as _load_csv
        result = _load_csv(
            file_path=file_path, preview_rows=preview_rows,
            max_preview_rows=settings.max_rows_display, data_dir=settings.data_dir,
        )
        log_action("AI_AGENT", "load_csv", f"Loaded {result['rows']} rows from {file_path}")
        return json.dumps(result, indent=2, default=str)

    except FileNotFoundError:
        return json.dumps({
            "error": f"File not found: {file_path}",
            "working_directory": str(Path.cwd()),
            "data_directory": str(Path(settings.data_dir)),
            "suggestions": [
                "1. Use find_files() to search for the file",
                "2. Use get_working_directory() to see available files",
                "3. Use stage_file() to copy the file to the data directory",
            ],
        }, indent=2)
    except Exception as e:
        return json.dumps({
            "error": str(e),
            "file_path": file_path,
            "suggestions": [
                "Use find_files() to search for the file",
                "Use get_working_directory() to check paths",
            ]
        })


@mcp.tool()
def load_json(file_path: str, preview_rows: int = 5) -> str:
    """
    Load a JSON file (array or object) and return a preview.

    Args:
        file_path: Path to the JSON file.
        preview_rows: Number of rows to preview (max 10).

    Returns:
        JSON with schema info and sample data.
    """
    try:
        from databridge.ingestion import load_json as _load_json
        result = _load_json(
            file_path=file_path, preview_rows=preview_rows,
            max_preview_rows=settings.max_rows_display,
        )
        log_action("AI_AGENT", "load_json", f"Loaded {result['rows']} records from {file_path}")
        return json.dumps(result, indent=2, default=str)

    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def query_database(connection_string: str, query: str, preview_rows: int = 10) -> str:
    """
    Execute a SQL query and return results.

    Args:
        connection_string: SQLAlchemy connection string (e.g., 'sqlite:///data.db').
        query: SQL SELECT query to execute.
        preview_rows: Maximum rows to return (max 10).

    Returns:
        JSON with query results and metadata.
    """
    try:
        from databridge.ingestion import query_database as _query_db
        result = _query_db(
            connection_string=connection_string, query=query,
            preview_rows=preview_rows, max_preview_rows=settings.max_rows_display,
        )
        log_action("AI_AGENT", "query_database", f"Query returned {result['rows_returned']} rows")
        return json.dumps(result, indent=2, default=str)

    except (ImportError, ValueError) as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 2: Data Profiling Tools
# =============================================================================

@mcp.tool()
def profile_data(source_path: str) -> str:
    """
    Analyze data structure and quality. Identifies table type and anomalies.

    Args:
        source_path: Path to CSV file to profile.

    Returns:
        JSON with profiling statistics including structure type, cardinality, and data quality metrics.
    """
    try:
        from databridge.profiler import profile_data as _profile
        summary = _profile(source_path=source_path)
        log_action("AI_AGENT", "profile_data", f"Profiled {source_path}")
        return json.dumps(summary, indent=2, default=str)

    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def detect_schema_drift(source_a_path: str, source_b_path: str) -> str:
    """
    Compare schemas between two CSV files to detect drift.

    Args:
        source_a_path: Path to first CSV (baseline).
        source_b_path: Path to second CSV (target).

    Returns:
        JSON with schema differences including added, removed, and type-changed columns.
    """
    try:
        from databridge.profiler import detect_schema_drift as _detect_drift
        result = _detect_drift(source_a_path=source_a_path, source_b_path=source_b_path)
        log_action("AI_AGENT", "detect_schema_drift", f"Compared schemas")
        return json.dumps(result, indent=2)

    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 3: Hashing & Comparison Engine
# =============================================================================

@mcp.tool()
def compare_hashes(
    source_a_path: str,
    source_b_path: str,
    key_columns: str,
    compare_columns: str = ""
) -> str:
    """
    Compare two CSV sources by hashing rows to identify orphans and conflicts.

    Args:
        source_a_path: Path to the first CSV file (source of truth).
        source_b_path: Path to the second CSV file (target).
        key_columns: Comma-separated column names that uniquely identify a row.
        compare_columns: Optional comma-separated columns to check for conflicts. Defaults to all non-key columns.

    Returns:
        JSON statistical summary with orphan and conflict counts (no raw data).
    """
    try:
        from databridge.reconciler import compare_hashes as _compare
        summary = _compare(
            source_a_path=source_a_path, source_b_path=source_b_path,
            key_columns=key_columns, compare_columns=compare_columns,
        )
        log_action("AI_AGENT", "compare_hashes",
                   f"Compared {summary['source_a']['total_rows']} vs {summary['source_b']['total_rows']} rows")
        return json.dumps(summary, indent=2)

    except ValueError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def get_orphan_details(
    source_a_path: str,
    source_b_path: str,
    key_columns: str,
    orphan_source: str = "both",
    limit: int = 10
) -> str:
    """
    Retrieve details of orphan records (records in one source but not the other).

    Args:
        source_a_path: Path to the first CSV file.
        source_b_path: Path to the second CSV file.
        key_columns: Comma-separated column names that uniquely identify a row.
        orphan_source: Which orphans to return: 'a', 'b', or 'both'.
        limit: Maximum orphans to return per source (max 10).

    Returns:
        JSON with orphan record details (limited to context sensitivity rules).
    """
    try:
        from databridge.reconciler import get_orphan_details as _get_orphans
        limit = min(limit, settings.max_rows_display)
        result = _get_orphans(
            source_a_path=source_a_path, source_b_path=source_b_path,
            key_columns=key_columns, orphan_source=orphan_source, limit=limit,
        )
        log_action("AI_AGENT", "get_orphan_details", f"Retrieved orphan details")
        return json.dumps(result, indent=2, default=str)

    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def get_conflict_details(
    source_a_path: str,
    source_b_path: str,
    key_columns: str,
    compare_columns: str = "",
    limit: int = 10
) -> str:
    """
    Retrieve details of conflicting records (same key, different values).

    Args:
        source_a_path: Path to the first CSV file.
        source_b_path: Path to the second CSV file.
        key_columns: Comma-separated column names that uniquely identify a row.
        compare_columns: Optional columns to compare. Defaults to all non-key columns.
        limit: Maximum conflicts to return (max 10).

    Returns:
        JSON with conflict details showing both versions side-by-side.
    """
    try:
        from databridge.reconciler import get_conflict_details as _get_conflicts
        limit = min(limit, settings.max_rows_display)
        result = _get_conflicts(
            source_a_path=source_a_path, source_b_path=source_b_path,
            key_columns=key_columns, compare_columns=compare_columns, limit=limit,
        )
        log_action("AI_AGENT", "get_conflict_details", f"Retrieved {result['showing']} conflict details")
        return json.dumps(result, indent=2, default=str)

    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 4: Fuzzy Matching Tools
# =============================================================================

@mcp.tool()
def fuzzy_match_columns(
    source_a_path: str,
    source_b_path: str,
    column_a: str,
    column_b: str,
    threshold: int = 80,
    limit: int = 10
) -> str:
    """
    Find fuzzy matches between two columns using RapidFuzz.

    Args:
        source_a_path: Path to the first CSV file.
        source_b_path: Path to the second CSV file.
        column_a: Column name in source A to match.
        column_b: Column name in source B to match against.
        threshold: Minimum similarity score (0-100). Default 80.
        limit: Maximum matches to return (max 10).

    Returns:
        JSON with fuzzy match results including similarity scores.
    """
    try:
        from databridge.reconciler import fuzzy_match_columns as _fuzzy_match
        limit = min(limit, settings.max_rows_display)
        result = _fuzzy_match(
            source_a_path=source_a_path, source_b_path=source_b_path,
            column_a=column_a, column_b=column_b,
            threshold=threshold, limit=limit,
        )
        log_action("AI_AGENT", "fuzzy_match_columns", f"Found {result['total_matches']} fuzzy matches")
        return json.dumps(result, indent=2)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def fuzzy_deduplicate(
    source_path: str,
    column: str,
    threshold: int = 90,
    limit: int = 10
) -> str:
    """
    Find potential duplicate values within a single column using fuzzy matching.

    Args:
        source_path: Path to the CSV file.
        column: Column name to check for duplicates.
        threshold: Minimum similarity score (0-100). Default 90.
        limit: Maximum duplicate groups to return (max 10).

    Returns:
        JSON with potential duplicate groups.
    """
    try:
        from databridge.reconciler import fuzzy_deduplicate as _fuzzy_dedup
        limit = min(limit, settings.max_rows_display)
        result = _fuzzy_dedup(
            source_path=source_path, column=column,
            threshold=threshold, limit=limit,
        )
        log_action("AI_AGENT", "fuzzy_deduplicate", f"Found {result['total_groups']} duplicate groups")
        return json.dumps(result, indent=2)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 5: PDF/OCR Tools
# =============================================================================

@mcp.tool()
def extract_text_from_pdf(file_path: str, pages: str = "all") -> str:
    """
    Extract text content from a PDF file.

    Args:
        file_path: Path to the PDF file.
        pages: Page numbers to extract ('all', or '1,2,3', or '1-5').

    Returns:
        JSON with extracted text per page.
    """
    try:
        from databridge.ingestion import extract_pdf_text as _extract_pdf
        result = _extract_pdf(file_path=file_path, pages=pages)
        log_action("AI_AGENT", "extract_text_from_pdf", f"Extracted {result['pages_extracted']} pages")
        return json.dumps(result, indent=2)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def ocr_image(file_path: str, language: str = "eng") -> str:
    """
    Extract text from an image using OCR (Tesseract).

    Args:
        file_path: Path to the image file (PNG, JPG, etc.).
        language: Tesseract language code (default 'eng').

    Returns:
        JSON with extracted text.
    """
    try:
        from databridge.ingestion import ocr_image as _ocr
        result = _ocr(
            file_path=file_path, language=language,
            tesseract_path=getattr(settings, 'tesseract_path', None),
        )
        log_action("AI_AGENT", "ocr_image", f"OCR extracted {result['character_count']} chars")
        return json.dumps(result, indent=2)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def parse_table_from_text(text: str, delimiter: str = "auto") -> str:
    """
    Attempt to parse tabular data from extracted text.

    Args:
        text: Raw text containing tabular data.
        delimiter: Column delimiter ('auto', 'tab', 'space', 'pipe', or custom).

    Returns:
        JSON with parsed table data.
    """
    try:
        from databridge.ingestion import parse_table_from_text as _parse_table
        result = _parse_table(
            text=text, delimiter=delimiter,
            max_preview_rows=settings.max_rows_display,
        )
        row_count = result.get("row_count", 1)
        log_action("AI_AGENT", "parse_table_from_text", f"Parsed {row_count} rows")
        return json.dumps(result, indent=2)

    except ValueError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 6: Workflow Management
# =============================================================================

from databridge.orchestrator import (
    save_workflow_step as _save_step,
    get_workflow as _get_workflow,
    clear_workflow as _clear_workflow,
    get_audit_log as _get_audit_log,
)


@mcp.tool()
def save_workflow_step(step_name: str, step_type: str, parameters: str) -> str:
    """Save a reconciliation step to the workflow recipe.

    Args:
        step_name: Descriptive name for this step.
        step_type: Type of operation (e.g., 'compare_hashes', 'fuzzy_match', 'transform').
        parameters: JSON string of parameters used for this step.

    Returns:
        Confirmation with updated workflow summary.
    """
    try:
        result = _save_step(settings.workflow_file, step_name, step_type, parameters)
        log_action("AI_AGENT", "save_workflow_step", f"Added step: {step_name}")
        return json.dumps(result)
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def get_workflow() -> str:
    """Retrieve the current workflow recipe.

    Returns:
        JSON with all workflow steps.
    """
    try:
        return json.dumps(_get_workflow(settings.workflow_file), indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def clear_workflow() -> str:
    """Clear all steps from the current workflow.

    Returns:
        Confirmation message.
    """
    try:
        result = _clear_workflow(settings.workflow_file)
        log_action("AI_AGENT", "clear_workflow", "Workflow cleared")
        return json.dumps(result)
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def get_audit_log(limit: int = 10) -> str:
    """Retrieve recent entries from the audit trail.

    Args:
        limit: Maximum entries to return (max 10).

    Returns:
        JSON with recent audit entries.
    """
    try:
        return json.dumps(
            _get_audit_log(settings.audit_log, limit, settings.max_rows_display),
            indent=2,
        )
    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 7: Data Transformation Tools
# =============================================================================

from databridge.reconciler import transform_column as _transform_column


@mcp.tool()
def transform_column(
    source_path: str,
    column: str,
    operation: str,
    output_path: str = ""
) -> str:
    """Apply a transformation to a column and optionally save the result.

    Args:
        source_path: Path to the CSV file.
        column: Column name to transform.
        operation: Transformation operation ('upper', 'lower', 'strip', 'trim_spaces', 'remove_special').
        output_path: Optional path to save transformed file. If empty, returns preview only.

    Returns:
        JSON with transformation preview and status.
    """
    try:
        result = _transform_column(source_path, column, operation, output_path)

        # Enhanced: Add character-level diff analysis for transformations
        try:
            from databridge.reconciler import diff_values_paired
            before = result["preview"]["before"]
            after = result["preview"]["after"]
            paired_diffs = diff_values_paired(before, after)
            result["preview"]["diffs"] = [
                {
                    "index": d.index,
                    "before": d.before,
                    "after": d.after,
                    "similarity": round(d.similarity, 4),
                    "changes": [
                        {"op": op.operation, "from": op.a_content, "to": op.b_content}
                        for op in d.opcodes if op.operation != "equal"
                    ],
                }
                for d in paired_diffs
            ]
        except ImportError:
            pass

        if output_path:
            log_action("AI_AGENT", "transform_column", f"Transformed {column} and saved to {output_path}")

        return json.dumps(result, indent=2)

    except (ValueError, Exception) as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def merge_sources(
    source_a_path: str,
    source_b_path: str,
    key_columns: str,
    merge_type: str = "inner",
    output_path: str = ""
) -> str:
    """
    Merge two CSV sources on key columns.

    Args:
        source_a_path: Path to the first CSV file.
        source_b_path: Path to the second CSV file.
        key_columns: Comma-separated column names to join on.
        merge_type: Type of merge ('inner', 'left', 'right', 'outer').
        output_path: Optional path to save merged file.

    Returns:
        JSON with merge statistics and preview.
    """
    try:
        from databridge.reconciler import merge_sources as _merge
        result = _merge(
            source_a_path=source_a_path, source_b_path=source_b_path,
            key_columns=key_columns, merge_type=merge_type,
            output_path=output_path, max_preview_rows=settings.max_rows_display,
        )
        if output_path:
            log_action("AI_AGENT", "merge_sources", f"Merged to {output_path}")
        return json.dumps(result, indent=2, default=str)

    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Phase 8: Living Documentation
# =============================================================================

@mcp.tool()
def update_manifest() -> str:
    """
    Regenerate the MANIFEST.md documentation from tool docstrings.

    Returns:
        Confirmation message with tool count.
    """
    try:
        tools = mcp._tool_manager._tools

        manifest = f"""# DataBridge AI - Tool Manifest

> Auto-generated documentation for all MCP tools.
> Last updated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

---

## Overview

DataBridge AI provides **{len(tools)} tools** for data reconciliation:

DataShield positioning: DataShield is an offline pre-AI masking layer.
It masks client data before AI processing and preserves table-purpose
context so post-processed data remains useful for modeling workflows.

| Category | Tools |
|----------|-------|
| Data Loading | load_csv, load_json, query_database |
| Profiling | profile_data, detect_schema_drift |
| Comparison | compare_hashes, get_orphan_details, get_conflict_details |
| Fuzzy Matching | fuzzy_match_columns, fuzzy_deduplicate |
| PDF/OCR | extract_text_from_pdf, ocr_image, parse_table_from_text |
| Workflow | save_workflow_step, get_workflow, clear_workflow, get_audit_log |
| Transform | transform_column, merge_sources |
| Documentation | update_manifest |

---

## Tool Reference

"""
        for name, tool in tools.items():
            doc = tool.fn.__doc__ or "No description available."
            manifest += f"### `{name}`\n\n{doc.strip()}\n\n---\n\n"

        manifest_path = Path("docs/MANIFEST.md")
        manifest_path.parent.mkdir(parents=True, exist_ok=True)

        with open(manifest_path, "w", encoding="utf-8") as f:
            f.write(manifest)

        log_action("AI_AGENT", "update_manifest", f"Updated with {len(tools)} tools")

        return json.dumps({
            "status": "success",
            "tools_documented": len(tools),
            "manifest_path": str(manifest_path)
        })

    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# License Status Tool
# =============================================================================

@mcp.tool()
def get_license_status() -> str:
    """
    Get the current license status and available features.

    Returns information about:
    - Current license tier (CE, PRO, ENTERPRISE, or FULL)
    - License validation status
    - Available modules at current tier
    - Upgrade information for additional features

    Returns:
        JSON with license status and feature availability.
    """
    if LICENSE_MANAGER:
        status = LICENSE_MANAGER.get_status()
        status['version'] = __version__
        return json.dumps(status, indent=2)

    # No license manager = full version
    return json.dumps({
        "version": __version__,
        "tier": "FULL",
        "tier_level": 2,
        "is_valid": True,
        "message": "Running full version (no license restrictions)",
        "available_modules": "all",
    }, indent=2)


# =============================================================================
# Phase 7B: Local SQL Tools (DuckDB acceleration)
# =============================================================================

@mcp.tool()
def query_local(sql: str, register_files: str = "{}") -> str:
    """
    Execute SQL against local CSV/Parquet/JSON files using DuckDB.

    DuckDB reads files natively in SQL:
      SELECT * FROM read_csv_auto('path/to/data.csv') WHERE amount > 1000
      SELECT * FROM read_parquet('data.parquet')
      SELECT * FROM read_json_auto('data.json')

    Args:
        sql: SQL query (DuckDB syntax). Use read_csv_auto(), read_parquet(), read_json_auto() for files.
        register_files: Optional JSON dict mapping table names to file paths for convenience.

    Returns:
        JSON with query results (column names, row count, preview data).

    Examples:
        query_local("SELECT department, AVG(salary) FROM read_csv_auto('employees.csv') GROUP BY 1")
        query_local("SELECT * FROM orders WHERE total > 1000", register_files='{"orders": "data/orders.csv"}')
    """
    try:
        from databridge.connectors import query_local as _query_local
        files = json.loads(register_files) if register_files and register_files != "{}" else None
        result = _query_local(
            sql=sql, register_files=files,
            max_preview_rows=settings.max_rows_display,
        )
        log_action("AI_AGENT", "query_local", f"DuckDB query returned {result['rows_returned']} rows")
        return json.dumps(result, indent=2, default=str)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e), "sql": sql})


@mcp.tool()
def register_local_table(file_path: str, table_name: str) -> str:
    """
    Register a local file as a named table in DuckDB for repeated SQL access.

    The table persists for the session — no need to use read_csv_auto() in every query.
    Supports CSV, Parquet, and JSON files (auto-detected by extension).

    Args:
        file_path: Path to CSV, Parquet, or JSON file.
        table_name: SQL table name to register.

    Returns:
        JSON confirmation with row count and column schema.
    """
    try:
        from databridge.connectors import register_local_table as _register
        info = _register(file_path=file_path, table_name=table_name)
        log_action("AI_AGENT", "register_local_table",
                   f"Registered {file_path} as '{table_name}' ({info.get('row_count', '?')} rows)")
        return json.dumps(info, indent=2, default=str)

    except (ImportError, FileNotFoundError) as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def list_local_tables() -> str:
    """
    List all tables registered in the local DuckDB instance.

    Returns:
        JSON with table names, row counts, and column schemas.
    """
    try:
        from databridge.connectors import list_local_tables as _list_tables
        result = _list_tables()
        log_action("AI_AGENT", "list_local_tables", f"Listed {result['table_count']} DuckDB tables")
        return json.dumps(result, indent=2, default=str)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


@mcp.tool()
def export_to_parquet(sql_or_table: str, output_path: str) -> str:
    """
    Export a DuckDB query result or registered table to Parquet format.

    Parquet is columnar, compressed, and ideal for Snowflake bulk loading (3-5x faster than CSV).

    Args:
        sql_or_table: SQL query string or a registered table name.
        output_path: Path for the output .parquet file.

    Returns:
        JSON with output path, row count, and file size.
    """
    try:
        from databridge.connectors import export_to_parquet as _export
        info = _export(sql_or_table=sql_or_table, output_path=output_path)
        log_action("AI_AGENT", "export_to_parquet",
                   f"Exported {info.get('row_count', '?')} rows to {output_path} ({info.get('file_size_mb', '?')} MB)")
        return json.dumps(info, indent=2, default=str)

    except ImportError as e:
        return json.dumps({"error": str(e)})
    except Exception as e:
        return json.dumps({"error": str(e)})


# =============================================================================
# Vanna RAG + PydanticAI — moved to src/plugins/admin/ (Phase 27, Rec 5)
# These 5 tools are now registered via plugin.json as ENTERPRISE tier.
# =============================================================================


# =============================================================================
# Dynamic Tool Registry — Gateway Mode for Cross-LLM Compatibility
# =============================================================================

try:
    from src.plugins.dynamic_registry import DynamicToolRegistry
except ImportError:
    from plugins.dynamic_registry import DynamicToolRegistry

_tool_mode = os.environ.get("DATABRIDGE_TOOL_MODE", "slim")
_dynamic_registry = DynamicToolRegistry(mcp, mode=_tool_mode)

# Capture metadata for core tools (registered above, before plugins)
_dynamic_registry.set_current_plugin("core", "CE")
_dynamic_registry.capture_tool_metadata()


# =============================================================================
# Plugin Loading — Dynamic registration of all external modules (Phase 9+)
# =============================================================================

try:
    try:
        from src.plugins.loader import load_all_plugins
    except ImportError:
        from plugins.loader import load_all_plugins

    _src_dir = Path(__file__).parent
    _project_root = _src_dir.parent

    plugin_results = load_all_plugins(
        mcp,
        settings,
        plugin_dirs=[
            _src_dir,                              # src/ (built-in modules)
            _project_root / "plugins",             # plugins/ (community)
            _project_root / "private_plugins",     # private_plugins/ (proprietary)
        ],
        license_manager=LICENSE_MANAGER,           # None = full mode (no restrictions)
        context={"dynamic_registry": _dynamic_registry},
    )

    _loaded = 0
    _skipped = 0
    _failed = 0
    for _name, _result in plugin_results.items():
        if _result.get("loaded"):
            log_action("SYSTEM", f"{_name}_init", f"{_name} registered ({_result.get('tools', '?')} tools)")
            _loaded += 1
        elif _result.get("skipped"):
            print(f"[License] {_name} requires {_result.get('tier', 'PRO')} license - skipped")
            _skipped += 1
        else:
            print(f"Warning: {_name} not loaded: {_result.get('error', 'unknown')}")
            _failed += 1

    print(f"[Plugins] Loaded {_loaded} plugins, {_skipped} skipped (license), {_failed} failed")

except Exception as _plugin_err:
    print(f"[Plugins] Plugin loader failed: {_plugin_err}")
    print("[Plugins] Falling back to core tools only")


# =============================================================================
# Gateway Tools — discover_tools, run_tool, list_domains, search_tools
# =============================================================================

@mcp.tool()
async def discover_tools(domain: str = "", query: str = "") -> str:
    """Discover available tools by domain or search query.

    Use list_domains() first to see all domains, then discover_tools(domain='blce')
    to list tools in that domain. Or use query to search by keyword.

    Args:
        domain: Domain name to list tools for (e.g. 'blce', 'catalog', 'wright').
        query: Search keyword to find tools across all domains.

    Returns:
        JSON array of matching tools (or domains if neither arg provided).
    """
    if query:
        return json.dumps(
            _annotate_tool_list_with_trust_preflight(
                _annotate_tool_list_with_trust_lane(_dynamic_registry.search_tools(query))
            ),
            indent=2,
        )
    if domain:
        return json.dumps(
            _annotate_tool_list_with_trust_preflight(
                _annotate_tool_list_with_trust_lane(_dynamic_registry.get_tools_by_domain(domain))
            ),
            indent=2,
        )
    return json.dumps(_dynamic_registry.get_all_domains(), indent=2)


@mcp.tool()
async def run_tool(
    tool_name: str,
    arguments: str = "{}",
    trust_attestation: str = "",
) -> str:
    """Execute any DataBridge tool by name.

    Use discover_tools() first to find the tool name and its parameters.
    Arguments should be a JSON string of the tool's parameters.

    Args:
        tool_name: The exact name of the tool to execute (e.g. 'blce_parse_sql').
        arguments: JSON string of tool parameters (e.g. '{"sql": "SELECT 1"}').

    Returns:
        The tool's result as a string.
    """
    args = json.loads(arguments)

    lane = _resolve_trust_lane_for_tool(tool_name)
    if lane:
        resolved_attestation = _coerce_trust_attestation_str(
            explicit_attestation=trust_attestation,
            arguments_dict=args,
        )
        trust_error = _enforce_trust_for_ai_tool(
            tool_name=tool_name,
            lane=lane,
            trust_attestation=resolved_attestation,
        )
        if trust_error:
            return json.dumps(trust_error)

    result = await _dynamic_registry.execute_tool(tool_name, args)
    return str(result)


@mcp.tool()
async def list_domains() -> str:
    """List all tool domains with descriptions and tool counts.

    Returns a JSON array of domains. Use discover_tools(domain='<name>') to see
    tools within a specific domain.

    Returns:
        JSON array of {domain, tool_count, description} objects.
    """
    return json.dumps(_dynamic_registry.get_all_domains(), indent=2)


@mcp.tool()
async def search_tools(query: str) -> str:
    """Search for tools by keyword across all domains.

    Searches tool names, descriptions, and domains for matches.

    Args:
        query: Search keyword (e.g. 'parse', 'sql', 'hierarchy').

    Returns:
        JSON array of matching tools with name, description, domain, and tier.
    """
    return json.dumps(
        _annotate_tool_list_with_trust_preflight(
            _annotate_tool_list_with_trust_lane(_dynamic_registry.search_tools(query))
        ),
        indent=2,
    )


# Capture gateway tool metadata, then finalize (hide tools in dynamic mode)
_dynamic_registry.set_current_plugin("gateway", "CE")
_dynamic_registry.capture_tool_metadata()
_dynamic_registry.finalize()

# Wire planner with full registry + purpose context
try:
    from src.agents.unified import set_planner_context
    set_planner_context(
        tool_registry=_dynamic_registry,
        purpose_path=str(Path(__file__).parent.parent / "purpose.md"),
    )
    print("[Planner] Context set: tool registry + purpose.md")
except Exception as _planner_err:
    print(f"[Planner] Context wiring failed: {_planner_err}")

if _tool_mode == "dynamic":
    print(f"[Gateway] Mode: dynamic — {_dynamic_registry.visible_count} visible, "
          f"{_dynamic_registry.hidden_count} hidden, {_dynamic_registry.total_count} total")


# =============================================================================
# Tenant Middleware (Phase 2A) — injects TenantContext before every tool call
# =============================================================================

try:
    from src.tenancy import TenantRegistry
    from src.tenancy.middleware import TenantMiddleware

    _tenant_registry = TenantRegistry(data_dir=Path(settings.data_dir) / "tenants")
    mcp.add_middleware(TenantMiddleware(registry=_tenant_registry))
    print("[Tenancy] Tenant middleware registered")
except Exception as _tenant_err:
    print(f"[Tenancy] Middleware not loaded: {_tenant_err}")


# =============================================================================
# Entry Point
# =============================================================================

if __name__ == "__main__":
    mcp.run()
