import asyncio
import os
import sys
import threading
import time
import traceback
from contextlib import asynccontextmanager
from datetime import datetime
from importlib import resources as importlib_resources
from typing import Any, Dict, List, Optional, Union

import pandas as pd
import psutil
from dateutil import parser
from fastapi import (
    Depends,
    FastAPI,
    Request,
    Response,
    WebSocket,
    WebSocketDisconnect,
    status,
)
from fastapi.concurrency import run_in_threadpool
from fastapi.logger import logger
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from google.protobuf.json_format import MessageToDict
from prometheus_client import Gauge, start_http_server
from pydantic import BaseModel

import feast
from feast import proto_json, utils
from feast.constants import DEFAULT_FEATURE_SERVER_REGISTRY_TTL
from feast.data_source import PushMode
from feast.errors import (
    FeastError,
    FeatureViewNotFoundException,
)
from feast.feast_object import FeastObject
from feast.permissions.action import WRITE, AuthzedAction
from feast.permissions.security_manager import assert_permissions
from feast.permissions.server.rest import inject_user_details
from feast.permissions.server.utils import (
    ServerType,
    init_auth_manager,
    init_security_manager,
    str_to_auth_manager_type,
)

# Define prometheus metrics
cpu_usage_gauge = Gauge(
    "feast_feature_server_cpu_usage", "CPU usage of the Feast feature server"
)
memory_usage_gauge = Gauge(
    "feast_feature_server_memory_usage", "Memory usage of the Feast feature server"
)


# TODO: deprecate this in favor of push features
class WriteToFeatureStoreRequest(BaseModel):
    feature_view_name: str
    df: dict
    allow_registry_cache: bool = True
    transform_on_write: bool = True


class PushFeaturesRequest(BaseModel):
    push_source_name: str
    df: dict
    allow_registry_cache: bool = True
    to: str = "online"
    transform_on_write: bool = True


class MaterializeRequest(BaseModel):
    start_ts: Optional[str] = None
    end_ts: Optional[str] = None
    feature_views: Optional[List[str]] = None
    disable_event_timestamp: bool = False


class MaterializeIncrementalRequest(BaseModel):
    end_ts: str
    feature_views: Optional[List[str]] = None


class GetOnlineFeaturesRequest(BaseModel):
    entities: Dict[str, List[Any]]
    feature_service: Optional[str] = None
    features: List[str] = []
    full_feature_names: bool = False


class GetOnlineDocumentsRequest(BaseModel):
    feature_service: Optional[str] = None
    features: List[str] = []
    full_feature_names: bool = False
    top_k: Optional[int] = None
    query: Optional[List[float]] = None
    query_string: Optional[str] = None
    api_version: Optional[int] = 1


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: List[ChatMessage]


class ReadDocumentRequest(BaseModel):
    file_path: str


class SaveDocumentRequest(BaseModel):
    file_path: str
    data: dict


async def _get_features(
    request: Union[GetOnlineFeaturesRequest, GetOnlineDocumentsRequest],
    store: "feast.FeatureStore",
):
    if request.feature_service:
        feature_service = await run_in_threadpool(
            store.get_feature_service, request.feature_service, allow_cache=True
        )
        assert_permissions(
            resource=feature_service, actions=[AuthzedAction.READ_ONLINE]
        )
        features = feature_service  # type: ignore
    else:
        all_feature_views, all_on_demand_feature_views = await run_in_threadpool(
            utils._get_feature_views_to_use,
            store.registry,
            store.project,
            request.features,
            allow_cache=True,
            hide_dummy_entity=False,
        )
        for feature_view in all_feature_views:
            assert_permissions(
                resource=feature_view, actions=[AuthzedAction.READ_ONLINE]
            )
        for od_feature_view in all_on_demand_feature_views:
            assert_permissions(
                resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE]
            )
        features = request.features  # type: ignore
    return features


def get_app(
    store: "feast.FeatureStore",
    registry_ttl_sec: int = DEFAULT_FEATURE_SERVER_REGISTRY_TTL,
):
    """
    Creates a FastAPI app that can be used to start a feature server.

    Args:
        store: The FeatureStore to use for serving features
        registry_ttl_sec: The TTL in seconds for the registry cache

    Returns:
        A FastAPI app

    Example:
        ```python
        from feast import FeatureStore

        store = FeatureStore(repo_path="feature_repo")
        app = get_app(store)
        ```

    The app provides the following endpoints:
    - `/get-online-features`: Get online features
    - `/retrieve-online-documents`: Retrieve online documents
    - `/push`: Push features to the feature store
    - `/write-to-online-store`: Write to the online store
    - `/health`: Health check
    - `/materialize`: Materialize features
    - `/materialize-incremental`: Materialize features incrementally
    - `/chat`: Chat UI
    - `/ws/chat`: WebSocket endpoint for chat
    MCP Support:
    - If MCP is enabled in feature server configuration, MCP endpoints will be added automatically
    """
    proto_json.patch()
    # Asynchronously refresh registry, notifying shutdown and canceling the active timer if the app is shutting down
    registry_proto = None
    shutting_down = False
    active_timer: Optional[threading.Timer] = None

    def stop_refresh():
        nonlocal shutting_down
        shutting_down = True
        if active_timer:
            active_timer.cancel()

    def async_refresh():
        if shutting_down:
            return

        store.refresh_registry()
        nonlocal registry_proto
        registry_proto = store.registry.proto()

        if registry_ttl_sec:
            nonlocal active_timer
            active_timer = threading.Timer(registry_ttl_sec, async_refresh)
            active_timer.start()

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        await store.initialize()
        async_refresh()
        yield
        stop_refresh()
        await store.close()

    app = FastAPI(lifespan=lifespan)

    @app.post(
        "/get-online-features",
        dependencies=[Depends(inject_user_details)],
    )
    async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, Any]:
        # Initialize parameters for FeatureStore.get_online_features(...) call
        features = await _get_features(request, store)

        read_params = dict(
            features=features,
            entity_rows=request.entities,
            full_feature_names=request.full_feature_names,
        )

        if store._get_provider().async_supported.online.read:
            response = await store.get_online_features_async(**read_params)  # type: ignore
        else:
            response = await run_in_threadpool(
                lambda: store.get_online_features(**read_params)  # type: ignore
            )

        # Convert the Protobuf object to JSON and return it
        response_dict = await run_in_threadpool(
            MessageToDict,
            response.proto,
            preserving_proto_field_name=True,
            float_precision=18,
        )
        return response_dict

    @app.post(
        "/retrieve-online-documents",
        dependencies=[Depends(inject_user_details)],
    )
    async def retrieve_online_documents(
        request: GetOnlineDocumentsRequest,
    ) -> Dict[str, Any]:
        logger.warning(
            "This endpoint is in alpha and will be moved to /get-online-features when stable."
        )
        # Initialize parameters for FeatureStore.retrieve_online_documents_v2(...) call
        features = await _get_features(request, store)

        read_params = dict(features=features, query=request.query, top_k=request.top_k)
        if request.api_version == 2 and request.query_string is not None:
            read_params["query_string"] = request.query_string

        if request.api_version == 2:
            response = await run_in_threadpool(
                lambda: store.retrieve_online_documents_v2(**read_params)  # type: ignore
            )
        else:
            response = await run_in_threadpool(
                lambda: store.retrieve_online_documents(**read_params)  # type: ignore
            )

        # Convert the Protobuf object to JSON and return it
        response_dict = await run_in_threadpool(
            MessageToDict,
            response.proto,
            preserving_proto_field_name=True,
            float_precision=18,
        )
        return response_dict

    @app.post("/push", dependencies=[Depends(inject_user_details)])
    async def push(request: PushFeaturesRequest) -> None:
        df = pd.DataFrame(request.df)
        actions = []
        if request.to == "offline":
            to = PushMode.OFFLINE
            actions = [AuthzedAction.WRITE_OFFLINE]
        elif request.to == "online":
            to = PushMode.ONLINE
            actions = [AuthzedAction.WRITE_ONLINE]
        elif request.to == "online_and_offline":
            to = PushMode.ONLINE_AND_OFFLINE
            actions = WRITE
        else:
            raise ValueError(
                f"{request.to} is not a supported push format. Please specify one of these ['online', 'offline', 'online_and_offline']."
            )

        from feast.data_source import PushSource

        all_fvs = store.list_feature_views(
            allow_cache=request.allow_registry_cache
        ) + store.list_stream_feature_views(allow_cache=request.allow_registry_cache)
        fvs_with_push_sources = {
            fv
            for fv in all_fvs
            if (
                fv.stream_source is not None
                and isinstance(fv.stream_source, PushSource)
                and fv.stream_source.name == request.push_source_name
            )
        }

        for feature_view in fvs_with_push_sources:
            assert_permissions(resource=feature_view, actions=actions)

        push_params = dict(
            push_source_name=request.push_source_name,
            df=df,
            allow_registry_cache=request.allow_registry_cache,
            to=to,
            transform_on_write=request.transform_on_write,
        )

        should_push_async = (
            store._get_provider().async_supported.online.write
            and to in [PushMode.ONLINE, PushMode.ONLINE_AND_OFFLINE]
        )
        if should_push_async:
            await store.push_async(**push_params)
        else:
            store.push(**push_params)

    async def _get_feast_object(
        feature_view_name: str, allow_registry_cache: bool
    ) -> FeastObject:
        # FIXME: this logic repeated at least 3 times in the codebase - should be centralized
        # in logging, in server and in feature_store (Python SDK)
        try:
            return await run_in_threadpool(
                store.get_feature_view,
                feature_view_name,
                allow_registry_cache=allow_registry_cache,
            )
        except FeatureViewNotFoundException:
            try:
                return await run_in_threadpool(
                    store.get_on_demand_feature_view,
                    feature_view_name,
                    allow_registry_cache=allow_registry_cache,
                )
            except FeatureViewNotFoundException:
                return await run_in_threadpool(
                    store.get_stream_feature_view,
                    feature_view_name,
                    allow_registry_cache=allow_registry_cache,
                )

    @app.post("/write-to-online-store", dependencies=[Depends(inject_user_details)])
    async def write_to_online_store(request: WriteToFeatureStoreRequest) -> None:
        df = pd.DataFrame(request.df)
        feature_view_name = request.feature_view_name
        allow_registry_cache = request.allow_registry_cache
        resource = await _get_feast_object(feature_view_name, allow_registry_cache)
        assert_permissions(resource=resource, actions=[AuthzedAction.WRITE_ONLINE])
        await run_in_threadpool(
            store.write_to_online_store,
            feature_view_name=feature_view_name,
            df=df,
            allow_registry_cache=allow_registry_cache,
            transform_on_write=request.transform_on_write,
        )

    @app.get("/health")
    async def health():
        return (
            Response(status_code=status.HTTP_200_OK)
            if registry_proto
            else Response(status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
        )

    @app.post("/chat")
    async def chat(request: ChatRequest):
        # Process the chat request
        # For now, just return dummy text
        return {"response": "This is a dummy response from the Feast feature server."}

    @app.post("/read-document")
    async def read_document_endpoint(request: ReadDocumentRequest):
        try:
            import os

            if not os.path.exists(request.file_path):
                return {"error": f"File not found: {request.file_path}"}

            with open(request.file_path, "r", encoding="utf-8") as file:
                content = file.read()

            return {"content": content, "file_path": request.file_path}
        except Exception as e:
            return {"error": str(e)}

    @app.post("/save-document")
    async def save_document_endpoint(request: SaveDocumentRequest):
        try:
            import json
            import os
            from pathlib import Path

            file_path = Path(request.file_path).resolve()
            if not str(file_path).startswith(os.getcwd()):
                return {"error": "Invalid file path"}

            base_name = file_path.stem
            labels_file = file_path.parent / f"{base_name}-labels.json"

            with open(labels_file, "w", encoding="utf-8") as file:
                json.dump(request.data, file, indent=2, ensure_ascii=False)

            return {"success": True, "saved_to": str(labels_file)}
        except Exception as e:
            return {"error": str(e)}

    @app.get("/chat")
    async def chat_ui():
        # Serve the chat UI
        static_dir_ref = importlib_resources.files(__spec__.parent) / "static/chat"  # type: ignore[name-defined, arg-type]
        with importlib_resources.as_file(static_dir_ref) as static_dir:
            with open(os.path.join(static_dir, "index.html")) as f:
                content = f.read()
        return Response(content=content, media_type="text/html")

    @app.post("/materialize", dependencies=[Depends(inject_user_details)])
    async def materialize(request: MaterializeRequest) -> None:
        for feature_view in request.feature_views or []:
            resource = await _get_feast_object(feature_view, True)
            assert_permissions(
                resource=resource,
                actions=[AuthzedAction.WRITE_ONLINE],
            )

        if request.disable_event_timestamp:
            # Query all available data and use current datetime as event timestamp
            now = datetime.now()
            start_date = datetime(
                1970, 1, 1
            )  # Beginning of time to capture all historical data
            end_date = now
        else:
            if not request.start_ts or not request.end_ts:
                raise ValueError(
                    "start_ts and end_ts are required when disable_event_timestamp is False"
                )
            start_date = utils.make_tzaware(parser.parse(request.start_ts))
            end_date = utils.make_tzaware(parser.parse(request.end_ts))

        await run_in_threadpool(
            store.materialize,
            start_date,
            end_date,
            request.feature_views,
            disable_event_timestamp=request.disable_event_timestamp,
        )

    @app.post("/materialize-incremental", dependencies=[Depends(inject_user_details)])
    async def materialize_incremental(request: MaterializeIncrementalRequest) -> None:
        for feature_view in request.feature_views or []:
            resource = await _get_feast_object(feature_view, True)
            assert_permissions(
                resource=resource,
                actions=[AuthzedAction.WRITE_ONLINE],
            )
        await run_in_threadpool(
            store.materialize_incremental,
            utils.make_tzaware(parser.parse(request.end_ts)),
            request.feature_views,
        )

    @app.exception_handler(Exception)
    async def rest_exception_handler(request: Request, exc: Exception):
        # Print the original exception on the server side
        logger.exception(traceback.format_exc())

        if isinstance(exc, FeastError):
            return JSONResponse(
                status_code=exc.http_status_code(),
                content=exc.to_error_detail(),
            )
        else:
            return JSONResponse(
                status_code=500,
                content=str(exc),
            )

    # Chat WebSocket connection manager
    class ConnectionManager:
        def __init__(self):
            self.active_connections: List[WebSocket] = []

        async def connect(self, websocket: WebSocket):
            await websocket.accept()
            self.active_connections.append(websocket)

        def disconnect(self, websocket: WebSocket):
            self.active_connections.remove(websocket)

        async def send_message(self, message: str, websocket: WebSocket):
            await websocket.send_text(message)

    manager = ConnectionManager()

    @app.websocket("/ws/chat")
    async def websocket_endpoint(websocket: WebSocket):
        await manager.connect(websocket)
        try:
            while True:
                message = await websocket.receive_text()
                # Process the received message (currently unused but kept for future implementation)
                # For now, just return dummy text
                response = f"You sent: '{message}'. This is a dummy response from the Feast feature server."

                # Stream the response word by word
                words = response.split()
                for word in words:
                    await manager.send_message(word + " ", websocket)
                    await asyncio.sleep(0.1)  # Add a small delay between words
        except WebSocketDisconnect:
            manager.disconnect(websocket)

    # Mount static files
    static_dir_ref = importlib_resources.files(__spec__.parent) / "static"  # type: ignore[name-defined, arg-type]
    with importlib_resources.as_file(static_dir_ref) as static_dir:
        app.mount("/static", StaticFiles(directory=static_dir), name="static")

    # Add MCP support if enabled in feature server configuration
    _add_mcp_support_if_enabled(app, store)

    return app


def _add_mcp_support_if_enabled(app, store: "feast.FeatureStore"):
    """Add MCP support to the FastAPI app if enabled in configuration."""
    try:
        # Check if MCP is enabled in feature server config
        if (
            store.config.feature_server
            and hasattr(store.config.feature_server, "type")
            and store.config.feature_server.type == "mcp"
            and getattr(store.config.feature_server, "mcp_enabled", False)
        ):
            from feast.infra.mcp_servers.mcp_server import add_mcp_support_to_app

            mcp_server = add_mcp_support_to_app(app, store, store.config.feature_server)

            if mcp_server:
                logger.info("MCP support has been enabled for the Feast feature server")
            else:
                logger.warning("MCP support was requested but could not be enabled")
        else:
            logger.debug("MCP support is not enabled in feature server configuration")
    except Exception as e:
        logger.error(f"Error checking/adding MCP support: {e}")
        # Don't fail the entire server if MCP fails to initialize


if sys.platform != "win32":
    import gunicorn.app.base

    class FeastServeApplication(gunicorn.app.base.BaseApplication):
        def __init__(self, store: "feast.FeatureStore", **options):
            self._app = get_app(
                store=store,
                registry_ttl_sec=options["registry_ttl_sec"],
            )
            self._options = options
            super().__init__()

        def load_config(self):
            for key, value in self._options.items():
                if key.lower() in self.cfg.settings and value is not None:
                    self.cfg.set(key.lower(), value)

            self.cfg.set("worker_class", "uvicorn_worker.UvicornWorker")

        def load(self):
            return self._app


def monitor_resources(self, interval: int = 5):
    """Function to monitor and update CPU and memory usage metrics."""
    logger.debug(f"Starting resource monitoring with interval {interval} seconds")
    p = psutil.Process()
    logger.debug(f"PID is {p.pid}")
    while True:
        with p.oneshot():
            cpu_usage = p.cpu_percent()
            memory_usage = p.memory_percent()
            logger.debug(f"CPU usage: {cpu_usage}%, Memory usage: {memory_usage}%")
            logger.debug(f"CPU usage: {cpu_usage}%, Memory usage: {memory_usage}%")
            cpu_usage_gauge.set(cpu_usage)
            memory_usage_gauge.set(memory_usage)
        time.sleep(interval)


def start_server(
    store: "feast.FeatureStore",
    host: str,
    port: int,
    no_access_log: bool,
    workers: int,
    keep_alive_timeout: int,
    registry_ttl_sec: int,
    tls_key_path: str,
    tls_cert_path: str,
    metrics: bool,
):
    if (tls_key_path and not tls_cert_path) or (not tls_key_path and tls_cert_path):
        raise ValueError(
            "Both key and cert file paths are required to start server in TLS mode."
        )
    if metrics:
        logger.info("Starting Prometheus Server")
        start_http_server(8000)

        logger.debug("Starting background thread to monitor CPU and memory usage")
        monitoring_thread = threading.Thread(
            target=monitor_resources, args=(5,), daemon=True
        )
        monitoring_thread.start()

    logger.debug("start_server called")
    auth_type = str_to_auth_manager_type(store.config.auth_config.type)
    logger.info(f"Auth type: {auth_type}")
    init_security_manager(auth_type=auth_type, fs=store)
    logger.debug("Security manager initialized successfully")
    init_auth_manager(
        auth_type=auth_type,
        server_type=ServerType.REST,
        auth_config=store.config.auth_config,
    )
    logger.debug("Auth manager initialized successfully")

    if sys.platform != "win32":
        options = {
            "bind": f"{host}:{port}",
            "accesslog": None if no_access_log else "-",
            "workers": workers,
            "keepalive": keep_alive_timeout,
            "registry_ttl_sec": registry_ttl_sec,
        }

        # Add SSL options if the paths exist
        if tls_key_path and tls_cert_path:
            options["keyfile"] = tls_key_path
            options["certfile"] = tls_cert_path
        FeastServeApplication(store=store, **options).run()
    else:
        import uvicorn

        app = get_app(store, registry_ttl_sec)
        if tls_key_path and tls_cert_path:
            uvicorn.run(
                app,
                host=host,
                port=port,
                access_log=(not no_access_log),
                ssl_keyfile=tls_key_path,
                ssl_certfile=tls_cert_path,
            )
        else:
            uvicorn.run(app, host=host, port=port, access_log=(not no_access_log))