from __future__ import annotations

"""Application factory for the Teradata MCP server.

High-level architecture:
- server.py is a thin entrypoint that parses CLI/env into a Settings object and
  calls create_mcp_app(settings), then runs the chosen transport.
- create_mcp_app builds a FastMCP instance, configures logging, adds middleware
  to capture per-request context (including stdio fast-path), sets up Teradata
  connections (and optional teradataml Feature Store), and registers tools,
  prompts and resources from both Python modules and YAML files.
- Tool handlers are plain Python functions named handle_<toolName> living under
  src/teradata_mcp_server/tools/*. They remain protocol-agnostic. At startup,
  we auto-wrap them with a small adapter so they appear as MCP tools with clean
  signatures. The adapter injects a DB connection and sets QueryBand from the
  request context when using HTTP.
"""
import inspect
import os
import re
from importlib.resources import files as pkg_files
from typing import Any

import yaml
from fastmcp import FastMCP
from fastmcp.prompts.prompt import TextContent, Message
from pydantic import Field

from teradata_mcp_server.config import Settings
from teradata_mcp_server import utils as config_utils
from teradata_mcp_server.utils import setup_logging, format_text_response, format_error_response
from teradata_mcp_server.middleware import RequestContextMiddleware
from teradata_mcp_server.tools.utils.queryband import build_queryband
from sqlalchemy.engine import Connection
from fastmcp.server.dependencies import get_context


def create_mcp_app(settings: Settings):
    """Create and configure the FastMCP app with middleware, tools, prompts, resources."""
    logger = setup_logging(settings.logging_level, settings.mcp_transport)

    # Load tool module loader via teradata tools package
    try:
        from teradata_mcp_server import tools as td
    except ImportError:
        import tools as td  # dev fallback

    mcp = FastMCP("teradata-mcp-server")

    # Profiles (load via utils to honor packaged + working-dir overrides)
    profile_name = settings.profile
    if not profile_name:
        logger.info("No profile specified, load all tools, prompts and resources.")
    config = config_utils.get_profile_config(profile_name)

    # Feature flags from profiles
    enableEFS = True if any(re.match(pattern, 'fs_*') for pattern in config.get('tool', [])) else False
    enableEVS = True if any(re.match(pattern, 'evs_*') for pattern in config.get('tool', [])) else False

    # Initialize TD connection and optional teradataml/EFS context
    # Pass settings object to TDConn instead of just connection_url
    tdconn = td.TDConn(settings=settings)
    fs_config = None
    if enableEFS:
        # Only import FeatureStoreConfig (which depends on tdfs4ds) when EFS tools are enabled
        try:
            from teradata_mcp_server.tools.fs.fs_utils import FeatureStoreConfig
            fs_config = FeatureStoreConfig()
            # teradataml is optional; warn if unavailable but keep EFS enabled
            try:
                import teradataml as tdml
                try:
                    tdml.create_context(tdsqlengine=tdconn.engine)
                except Exception as e:
                    logger.warning(f"Error creating teradataml context: {e}")
            except (AttributeError, ImportError, ModuleNotFoundError):
                logger.warning("teradataml not installed; EFS tools will operate without a teradataml context")
        except (AttributeError, ImportError, ModuleNotFoundError) as e:
            logger.warning(f"Feature Store module not available - disabling EFS functionality: {e}")
            enableEFS = False

    # EVS connection (optional)
    evs = None
    if len(os.getenv("VS_NAME", "").strip()) > 0:
        try:
            evs = td.get_evs()
            enableEVS = True
        except Exception as e:
            logger.error(f"Unable to establish connection to EVS, disabling: {e}")
            enableEVS = False

    # Middleware (auth + request context)
    from teradata_mcp_server.tools.auth_cache import SecureAuthCache
    auth_cache = SecureAuthCache(ttl_seconds=settings.auth_cache_ttl)

    def get_tdconn(recreate: bool = False):
        nonlocal tdconn, fs_config
        if recreate:
            tdconn = td.TDConn(settings=settings)
            if enableEFS:
                try:
                    import teradataml as tdml
                    fs_config = td.FeatureStoreConfig()
                    try:
                        tdml.create_context(tdsqlengine=tdconn.engine)
                    except Exception:
                        pass
                except Exception:
                    pass
        return tdconn

    middleware = RequestContextMiddleware(
        logger=logger,
        auth_cache=auth_cache,
        tdconn_supplier=get_tdconn,
        auth_mode=settings.auth_mode,
        transport=settings.mcp_transport,
    )
    mcp.add_middleware(middleware)

    # Adapters (inlined for simplicity)
    process_id = f"{os.uname().nodename}:{os.getpid()}"

    def execute_db_tool(tool, *args, **kwargs):
        """Execute a handler with a DB connection and MCP concerns.

        - Detects whether the handler expects a SQLAlchemy Connection or a raw
          DB-API connection and injects appropriately.
        - For HTTP transport, builds and sets Teradata QueryBand per request using
          the RequestContext captured by middleware.
        - Formats return values into FastMCP content and captures exceptions with
          context for easier debugging.
        """
        tool_name = kwargs.pop('tool_name', getattr(tool, '__name__', 'unknown_tool'))
        tdconn_local = get_tdconn()

        if not getattr(tdconn_local, "engine", None):
            logger.info("Reinitializing TDConn")
            tdconn_local = get_tdconn(recreate=True)

        sig = inspect.signature(tool)
        first_param = next(iter(sig.parameters.values()))
        ann = first_param.annotation
        use_sqla = inspect.isclass(ann) and issubclass(ann, Connection)

        try:
            if use_sqla:
                from sqlalchemy import text
                with tdconn_local.engine.connect() as conn:
                    # Always attempt to set QueryBand when a request context is present
                    ctx = get_context()
                    request_context = ctx.get_state("request_context") if ctx else None
                    if request_context is not None:
                        qb = build_queryband(
                            application=mcp.name,
                            profile=profile_name,
                            process_id=process_id,
                            tool_name=tool_name,
                            request_context=request_context,
                        )
                        try:
                            conn.execute(text(f"SET QUERY_BAND = '{qb}' FOR SESSION"))
                            logger.debug(f"QueryBand set: {qb}")
                            logger.debug(f"Tool request context: {request_context}")
                        except Exception as qb_error:
                            logger.debug(f"Could not set QueryBand: {qb_error}")
                            # If in Basic auth, do not run the tool without proxying
                            if str(getattr(request_context, "auth_scheme", "")).lower() == "basic":
                                return format_error_response(
                                    f"Cannot run tool '{tool_name}': failed to set QueryBand for Basic auth. Error: {qb_error}"
                                )
                    result = tool(conn, *args, **kwargs)
            else:
                raw = tdconn_local.engine.raw_connection()
                try:
                    # Always attempt to set QueryBand when a request context is present
                    ctx = get_context()
                    request_context = ctx.get_state("request_context") if ctx else None
                    if request_context is not None:
                        qb = build_queryband(
                            application=mcp.name,
                            profile=profile_name,
                            process_id=process_id,
                            tool_name=tool_name,
                            request_context=request_context,
                        )
                        try:
                            cursor = raw.cursor()
                            # Apply at session scope so it persists across statements
                            cursor.execute(f"SET QUERY_BAND = '{qb}' FOR SESSION")
                            cursor.close()
                            logger.debug(f"QueryBand set: {qb}")
                            logger.debug(f"Tool request context: {request_context}")
                        except Exception as qb_error:
                            logger.debug(f"Could not set QueryBand: {qb_error}")
                            if str(getattr(request_context, "auth_scheme", "")).lower() == "basic":
                                return format_error_response(
                                    f"Cannot run tool '{tool_name}': failed to set QueryBand for Basic auth. Error: {qb_error}"
                                )
                    result = tool(raw, *args, **kwargs)
                finally:
                    raw.close()
            return format_text_response(result)
        except Exception as e:
            logger.error(f"Error in execute_db_tool: {e}", exc_info=True, extra={"session_info": {"tool_name": tool_name}})
            return format_error_response(str(e))

    def make_tool_wrapper(func):
        """Create an MCP-facing wrapper for a handle_* function.

        - Removes internal parameters (conn, tool_name, fs_config) from the MCP
          signature while still injecting them into the underlying handler.
        - Preserves the handler's parameter names and types so MCP clients can
          render friendly forms.
        """
        sig = inspect.signature(func)
        inject_kwargs = {}
        removable = {"conn", "tool_name"}
        if "fs_config" in sig.parameters:
            inject_kwargs["fs_config"] = fs_config
            removable.add("fs_config")

        params = [
            p for name, p in sig.parameters.items()
            if name not in removable and p.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
        ]
        new_sig = sig.replace(parameters=params)

        # Preserve annotations for Pydantic schema generation
        annotations = {}
        for name, p in sig.parameters.items():
            if name in removable:
                continue
            if p.kind not in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY):
                continue
            if p.annotation is not inspect._empty:
                annotations[name] = p.annotation

        def _exec(*args, **kwargs):
            return execute_db_tool(func, **inject_kwargs, **kwargs)

        _exec.__name__ = getattr(func, "__name__", "wrapped_tool")
        _exec.__signature__ = new_sig
        if annotations:
            _exec.__annotations__ = annotations
        return _exec

    # Register code tools via module loader
    module_loader = td.initialize_module_loader(config)
    if module_loader:
        all_functions = module_loader.get_all_functions()
        for name, func in all_functions.items():
            if not (inspect.isfunction(func) and name.startswith("handle_")):
                continue
            tool_name = name[len("handle_"):]
            if not any(re.match(p, tool_name) for p in config.get('tool', [])):
                continue
            wrapped = make_tool_wrapper(func)
            mcp.tool(name=tool_name, description=wrapped.__doc__)(wrapped)
            logger.info(f"Created tool: {tool_name}")
    else:
        logger.warning("No module loader available, skipping code-defined tool registration")

    # Load YAML-defined tools/resources/prompts
    custom_object_files = [file for file in os.listdir() if file.endswith("_objects.yml")]
    if module_loader and profile_name:
        profile_yml_files = module_loader.get_required_yaml_paths()
        custom_object_files.extend(profile_yml_files)
        logger.info(f"Loading YAML files for profile '{profile_name}': {len(profile_yml_files)} files")
    else:
        tool_yml_resources = []
        tools_pkg_root = pkg_files("teradata_mcp_server").joinpath("tools")
        if tools_pkg_root.is_dir():
            for subpkg in tools_pkg_root.iterdir():
                if subpkg.is_dir():
                    for entry in subpkg.iterdir():
                        if entry.is_file() and entry.name.endswith('.yml'):
                            tool_yml_resources.append(entry)
        custom_object_files.extend(tool_yml_resources)
        logger.info(f"Loading all YAML files (no specific profile): {len(tool_yml_resources)} files")

    custom_objects: dict[str, Any] = {}
    custom_glossary: dict[str, Any] = {}
    for file in custom_object_files:
        try:
            if hasattr(file, "read_text"):
                text = file.read_text(encoding='utf-8')
            else:
                with open(file, encoding='utf-8', errors='replace') as f:
                    text = f.read()
            loaded = yaml.safe_load(text)
            if loaded:
                custom_objects.update(loaded)
        except Exception as e:
            logger.error(f"Failed to load YAML from {file}: {e}")

    # Prompt helpers
    def make_custom_prompt(prompt_name: str, prompt: str, desc: str, parameters: dict | None = None):
        if parameters is None or len(parameters) == 0:
            async def _dynamic_prompt():
                return Message(role="user", content=TextContent(type="text", text=prompt))
            _dynamic_prompt.__name__ = prompt_name
            return mcp.prompt(description=desc)(_dynamic_prompt)
        else:
            param_objects: list[inspect.Parameter] = []
            annotations: dict[str, Any] = {}
            for param_name, meta in parameters.items():
                meta = meta or {}
                type_hint_raw = meta.get("type_hint", str)
                if isinstance(type_hint_raw, str):
                    try:
                        type_hint = eval(type_hint_raw, {"str": str, "int": int, "float": float, "bool": bool})
                    except Exception:
                        type_hint = str
                else:
                    type_hint = type_hint_raw
                required = meta.get("required", True)
                desc_txt = meta.get("description", "")
                desc_txt += f" (type: {type_hint_raw})"
                if required and "default" not in meta:
                    default_value = Field(..., description=desc_txt)
                else:
                    default_value = Field(default=meta.get("default", None), description=desc_txt)
                param_objects.append(
                    inspect.Parameter(
                        param_name,
                        kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
                        default=default_value,
                        annotation=type_hint,
                    )
                )
                annotations[param_name] = type_hint
            sig = inspect.Signature(param_objects)
            async def _dynamic_prompt(**kwargs):
                missing = [name for name, meta in parameters.items() if (meta or {}).get("required", True) and name not in kwargs]
                if missing:
                    raise ValueError(f"Missing parameters: {missing}")
                formatted_prompt = prompt.format(**kwargs)
                return Message(role="user", content=TextContent(type="text", text=formatted_prompt))
            _dynamic_prompt.__signature__ = sig
            _dynamic_prompt.__annotations__ = annotations
            _dynamic_prompt.__name__ = prompt_name
            return mcp.prompt(description=desc)(_dynamic_prompt)

    def make_custom_query_tool(name, tool):
        param_defs = tool.get("parameters", {})
        parameters = []
        annotations = {}
        for param_name, p in param_defs.items():
            type_hint = p.get("type_hint", str)
            default = inspect.Parameter.empty if p.get("required", True) else p.get("default", None)
            parameters.append(
                inspect.Parameter(param_name, kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, default=default, annotation=type_hint)
            )
            annotations[param_name] = type_hint
        sig = inspect.Signature(parameters)
        async def _dynamic_tool(**kwargs):
            missing = [n for n in annotations if n not in kwargs]
            if missing:
                raise ValueError(f"Missing parameters: {missing}")
            return execute_db_tool(td.handle_base_readQuery, tool["sql"], tool_name=name, **kwargs)
        _dynamic_tool.__signature__ = sig
        _dynamic_tool.__annotations__ = annotations
        return mcp.tool(name=name, description=tool.get("description", ""))(_dynamic_tool)

    def generate_cube_query_tool(name, cube):
        """
        Generate a function to create aggregation SQL from a cube definition.

        :param cube: The cube definition
        :return: A SQL query string generator function taking dimensions and measures as comma-separated strings.
        """
        def _cube_query_tool(dimensions: str, measures: str, dim_filters: str, meas_filters: str, order_by: str, top: int) -> str:
            """
            Generate a SQL query string for the cube using the specified dimensions and measures.

            Args:
                dimensions (str): Comma-separated dimension names (keys in cube['dimensions']).
                measures (str): Comma-separated measure names (keys in cube['measures']).
                dim_filters (str): Filter SQL expressions on dimensions.
                meas_filters (str): Filter SQL expressions on computed measures.
                order_by (str): Order SQL expressions on selected dimensions and measures.
                top (int): Filters the top N results.

            Returns:
                str: The generated SQL query.
            """
            dim_list_raw = [d.strip() for d in dimensions.split(",") if d.strip()]
            mes_list_raw = [m.strip() for m in measures.split(",") if m.strip()]
            # Get dimension expressions from dictionary
            dim_list = ",\n  ".join([
                cube["dimensions"][d]["expression"] if d in cube["dimensions"] else d
                for d in dim_list_raw
            ])
            mes_lines = []
            for measure in mes_list_raw:
                mdef = cube["measures"].get(measure)
                if mdef is None:
                    raise ValueError(f"Measure '{measure}' not found in cube '{name}'.")
                expr = mdef["expression"]
                mes_lines.append(f"{expr} AS {measure}")
            meas_list = ",\n  ".join(mes_lines)
            top_clause = f"TOP {top}" if top else ""
            dim_comma = ",\n  " if dim_list.strip() else ""
            where_dim_clause = f"WHERE {dim_filters}" if dim_filters else ""
            where_meas_clause = f"WHERE {meas_filters}" if meas_filters else ""
            order_clause = f"ORDER BY {order_by}" if order_by else ""
            
            sql = (
                f"SELECT {top_clause} * from\n"
                "(SELECT\n"
                f"  {dim_list}{dim_comma}"
                f"  {meas_list}\n"
                "FROM (\n"
                f"{cube['sql'].strip()}\n"
                f"{where_dim_clause}"
                ") AS c\n"
                f"GROUP BY {', '.join(dim_list_raw)}"
                ") AS a\n"
                f"{where_meas_clause}"
                f"{order_clause}"
                ";"            
            )
            return sql
        return _cube_query_tool

    def make_custom_cube_tool(name, cube):
        async def _dynamic_tool(dimensions, measures, dim_filters="", meas_filters="", order_by="", top=None):
            # Accept dimensions and measures as comma-separated strings, parse to lists
            return execute_db_tool(
                td.util_base_dynamicQuery,
                sql_generator=generate_cube_query_tool(name, cube),
                dimensions=dimensions,
                measures=measures,
                dim_filters=dim_filters,
                meas_filters=meas_filters,
                order_by=order_by,
                top=top
            )
        _dynamic_tool.__name__ = 'get_cube_' + name
        # Build allowed values and definitions for dimensions and measures
        dim_lines = []
        for n, d in cube.get('dimensions', {}).items():
            dim_lines.append(f"    - {n}: {d.get('description', '')}")
        measure_lines = []
        for n, m in cube.get('measures', {}).items():
            measure_lines.append(f"    - {n}: {m.get('description', '')}")
        
        # Create example strings for documentation
        dim_examples = [f"{d} {e}" for d, e in zip(list(cube.get('dimensions', {}))[:2], ["= 'value'", "in ('X', 'Y', 'Z')"])]
        dim_example = ' AND '.join(dim_examples)
        
        meas_examples = [f"{m} {e}" for m, e in zip(list(cube.get('measures', {}))[:2], ["> 1000", "= 100"])]
        meas_example = ' AND '.join(meas_examples)
        
        order_examples = [f"{d} {e}" for d, e in zip(list(cube.get('dimensions', {}))[:2], [" ASC", " DESC"])]
        order_example = ' , '.join(order_examples)
        
        _dynamic_tool.__doc__ = f"""
        Tool to query the cube '{name}'.
        {cube.get('description', '')}

        Expected inputs:
            * dimensions (str): Comma-separated dimension names to group by. Allowed values for dimensions\n:
    {chr(10).join(dim_lines)}

            * measures (str): Comma-separated measure names to aggregate. Allowed values for measures:
    {chr(10).join(measure_lines)}

            * dim_filters (str): Filter expression to apply to dimensions. Valid dimension names are: [{', '.join(cube.get('dimensions', {}).keys())}], use valid SQL expressions, for example:
    "{dim_example}"
            * meas_filters (str): Filter expression to apply to computed measures. Valid measure names are: [{', '.join(cube.get('measures', {}).keys())}], use valid SQL expressions, for example:
    "{meas_example}"
            * order_by (str): Order expression on any selected dimensions and measures. Use SQL syntax, for example:
    "{order_example}"
            top (int): Limit the number of rows returned, use a positive integer.

        Returns:
            Query result as a formatted response.
        """
        return mcp.tool(description=_dynamic_tool.__doc__)(_dynamic_tool)

    # Register custom objects
    custom_terms: list[tuple[str, Any, str]] = []
    for name, obj in custom_objects.items():
        obj_type = obj.get("type")
        if obj_type == "tool" and any(re.match(pattern, name) for pattern in config.get('tool',[])):
            fn = make_custom_query_tool(name, obj)
            globals()[name] = fn
            logger.info(f"Created tool: {name}")
        elif obj_type == "prompt"  and any(re.match(pattern, name) for pattern in config.get('prompt',[])):
            fn = make_custom_prompt(name, obj["prompt"], obj.get("description", ""), obj.get("parameters", {}))
            globals()[name] = fn
            logger.info(f"Created prompt: {name}")
        elif obj_type == "cube"  and any(re.match(pattern, name) for pattern in config.get('tool',[])):
            fn = make_custom_cube_tool(name, obj)
            globals()[name] = fn
            logger.info(f"Created cube: {name}")
        elif obj_type == "glossary"  and any(re.match(pattern, name) for pattern in config.get('resource',[])):
            custom_glossary = {k: v for k, v in obj.items() if k != "type"}
            logger.info(f"Added custom glossary entries for: {name}.")
        else:
            logger.info(f"Type {obj_type if obj_type else ''} for custom object {name} is {'unknown' if obj_type else 'undefined'}.")

        for section in ("measures", "dimensions"):
            if section in obj and  any(re.match(pattern, name) for pattern in config.get('tool',[])):
                custom_terms.extend((term, details, name) for term, details in obj[section].items())

    # Enrich glossary
    for term, details, tool_name in custom_terms:
        term_key = term.strip()
        if term_key not in custom_glossary:
            custom_glossary[term_key] = {"definition": details.get("description"), "synonyms": [], "tools": [tool_name]}
        else:
            if "tools" not in custom_glossary[term_key]:
                custom_glossary[term_key]["tools"] = []
            if tool_name not in custom_glossary[term_key]["tools"]:
                custom_glossary[term_key]["tools"].append(tool_name)

    if custom_glossary:
        @mcp.resource("glossary://all")
        def get_glossary() -> dict:
            return custom_glossary

        @mcp.resource("glossary://definitions")
        def get_glossary_definitions() -> dict:
            return {term: details["definition"] for term, details in custom_glossary.items()}

        @mcp.resource("glossary://term/{term_name}")
        def get_glossary_term(term_name: str)  -> dict:
            term = custom_glossary.get(term_name)
            if term:
                return term
            else:
                return {"error": f"Glossary term not found: {term_name}"}

    # Return the configured app and some handles used by the entrypoint if needed
    return mcp, logger
