from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from guardrails import configure_logging
from guardrails_api.clients.cache_client import CacheClient
from guardrails_api.db.postgres_client import postgres_is_enabled

from rich.console import Console
from rich.rule import Rule
from typing import Optional
import importlib.util
import json
import os
import inspect


GR_ENV_FILE = os.environ.get("GR_ENV_FILE", None)
GR_CONFIG_FILE_PATH = os.environ.get("GR_CONFIG_FILE_PATH", None)
GR_MIDDLEWARE_FILE_PATH = os.environ.get("GR_MIDDLEWARE_FILE_PATH", None)
PORT = int(os.environ.get("PORT", 8000))


# Custom JSON encoder
class CustomJSONEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, set):
            return list(o)
        if callable(o):
            return str(o)
        return super().default(o)


def register_config(config: Optional[str] = None):
    default_config_file = os.path.join(os.getcwd(), "./config.py")
    config_file = config or default_config_file
    config_file_path = os.path.abspath(config_file)
    if os.path.isfile(config_file_path):
        spec = importlib.util.spec_from_file_location("config", config_file_path)
        if spec and spec.loader:
            config_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(config_module)

    return config_file_path


def register_middleware(*, middleware: Optional[str] = None, app: FastAPI):
    default_middleware_file = os.path.join(os.getcwd(), "./middleware.py")
    middleware_file = middleware or default_middleware_file
    middleware_file_path = os.path.abspath(middleware_file)
    if os.path.isfile(middleware_file_path):
        try:
            spec = importlib.util.spec_from_file_location(
                "middleware", middleware_file_path
            )
            if spec and spec.loader:
                middleware_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(middleware_module)

                exports = middleware_module.__dir__()
                for export_name in exports:
                    export = getattr(middleware_module, export_name)
                    is_middleware = (
                        inspect.isclass(export)
                        and issubclass(export, BaseHTTPMiddleware)
                        and export != BaseHTTPMiddleware
                    )
                    if is_middleware:
                        app.add_middleware(export)
        except Exception as e:
            raise RuntimeError(
                f"Failed to register middleware from {middleware_file_path}", e
            )


# Support for providing env vars as uvicorn does not support supplying args to create_app
# - Usage: GR_CONFIG_FILE_PATH=config.py GR_ENV_FILE=.env PORT=8080 uvicorn --factory 'guardrails_api.app:create_app' --host 0.0.0.0 --port $PORT --workers 2 --timeout-keep-alive 90
# - Usage: gunicorn -k uvicorn.workers.UvicornWorker --bind 0.0.0.0:$PORT --timeout=90 --workers=2 "guardrails_api.app:create_app(None, None, $PORT)"
def create_app(
    env: Optional[str] = GR_ENV_FILE,
    config: Optional[str] = GR_CONFIG_FILE_PATH,
    port: Optional[int] = PORT,
    *,
    middleware: Optional[str] = GR_MIDDLEWARE_FILE_PATH,
    env_override: bool = False,
):
    # used to print user-facing messages during server startup
    console = Console()

    # Load specified env file or default .env file if it exists
    env_file_path = os.path.abspath(env or os.path.join(os.getcwd(), "./.env"))
    if os.path.isfile(env_file_path):
        from dotenv import load_dotenv

        load_dotenv(env_file_path, override=env_override)

    set_port = port or PORT
    host = os.environ.get("HOST", "http://localhost")
    self_endpoint = os.environ.get("SELF_ENDPOINT") or f"{host}:{set_port}"
    os.environ["SELF_ENDPOINT"] = self_endpoint

    resolved_config_file_path = register_config(config)

    app = FastAPI(openapi_url="")

    # Add CORS middleware
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    register_middleware(middleware=middleware, app=app)

    guardrails_log_level = os.environ.get("GUARDRAILS_LOG_LEVEL", "INFO")
    configure_logging(log_level=guardrails_log_level)

    # if no pg_host is set, don't set up postgres
    if postgres_is_enabled():
        from guardrails_api.db.postgres_client import PostgresClient

        pg_client = PostgresClient()
        pg_client.initialize(app)

    cache_client = CacheClient()
    cache_client.initialize()

    from guardrails_api.api.root import router as root_router
    from guardrails_api.api.guards import router as guards_router
    from guardrails_api.clients.get_guard_client import get_guard_client

    app.include_router(root_router)
    app.include_router(guards_router)

    # Custom JSON encoder
    @app.exception_handler(ValueError)
    async def value_error_handler(request: Request, exc: ValueError):
        return JSONResponse(
            status_code=400,
            content={"message": str(exc)},
        )

    console.print(f"\n:rocket: Guardrails API is available at {self_endpoint}")
    console.print(
        f":book: Visit {self_endpoint}/docs to see available API endpoints.\n"
    )

    console.print(":green_circle: Active guards and OpenAI compatible endpoints:")

    guard_client = get_guard_client()
    guards = guard_client.get_guards()

    for g in guards:
        g_dict = g.to_dict()
        console.print(
            f"- Guard: [bold white]{g_dict.get('name')}[/bold white] {self_endpoint}/guards/{g_dict.get('name')}/openai/v1"
        )

    console.print("")
    console.print("Using the following configuration:")
    console.print(f"- Guardrails Log Level: {guardrails_log_level}")
    console.print(f"- Self Endpoint: {self_endpoint}")
    console.print(
        f"- Config File Path: {resolved_config_file_path} [Provided: {config}]"
    )
    console.print(
        Rule("[bold grey]Server Logs[/bold grey]", characters="=", style="white")
    )

    return app


if __name__ == "__main__":
    import uvicorn

    app = create_app()
    uvicorn.run(app, host="0.0.0.0", port=PORT)
