import fcntl
import logging
import os
import random
import time
from contextlib import asynccontextmanager
from typing import AsyncGenerator

import torch
import uvicorn
from fastapi import APIRouter, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse
from opentelemetry import _logs, trace
from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import Response
from starlette.types import Lifespan

from auth.SPAStaticFiles import SPAStaticFiles
from clients.telemetry.telemetry_client import TelemetryEventTypes, send_telemetry_event
from config.config import Config
from config.extra_features import extra_feature_config
from dependencies import (
    db_session_context,
    get_db_engine,
    get_db_session,
    get_keycloak_client,
    get_keycloak_settings,
    get_oauth_client,
    get_scorer_client,
)
from monitoring.audit_log_middleware import AuditLogMiddleware, setup_audit_logger
from repositories.system_task_repository import SystemTaskRepository
from routers.api_key_routes import api_keys_routes
from routers.auth_routes import auth_routes
from routers.chat_routes import app_chat_routes
from routers.health_routes import health_router
from routers.user_routes import user_identity_routes, user_management_routes
from routers.v1.agent_polling_routes import agent_polling_routes
from routers.v1.agentic_experiment_routes import agentic_experiment_routes
from routers.v1.agentic_notebook_routes import agentic_notebook_routes
from routers.v1.agentic_prompt_routes import agentic_prompt_routes
from routers.v1.chatbot_routes import chatbot_routes
from routers.v1.continuous_eval_routes import continuous_eval_routes
from routers.v1.demo_task_routes import demo_task_routes
from routers.v1.legacy_span_routes import span_routes
from routers.v1.llm_eval_routes import llm_eval_routes
from routers.v1.model_provider_routes import model_provider_routes
from routers.v1.notebook_routes import notebook_routes
from routers.v1.prompt_experiment_routes import prompt_experiment_routes
from routers.v1.rag_experiment_routes import rag_experiment_routes
from routers.v1.rag_notebook_routes import rag_notebook_routes
from routers.v1.rag_routes import rag_routes
from routers.v1.rag_setting_routes import rag_setting_routes
from routers.v1.secrets_routes import secrets_routes
from routers.v1.trace_api_routes import trace_api_routes
from routers.v1.transform_routes import transform_routes
from routers.v2.routers import (
    dataset_management_routes,
    engine_config_routes,
    feedback_routes,
    query_routes,
    rule_management_routes,
    system_management_routes,
    task_management_routes,
    validate_routes,
)
from routers.v2.tenant_signup_routes import tenant_signup_routes
from services.continuous_eval import (
    initialize_continuous_eval_queue_service,
    shutdown_continuous_eval_queue_service,
)
from services.currency import (
    initialize_currency_conversion_service,
    shutdown_currency_conversion_service,
)
from services.system_tasks_service import initialize_system_tasks
from services.task import (
    initialize_global_agent_polling_service,
    shutdown_global_agent_polling_service,
)
from services.trace_retention_service import (
    initialize_trace_retention_service,
    shutdown_trace_retention_service,
)
from utils import constants as constants
from utils import model_load
from utils.classifiers import get_device
from utils.utils import (
    get_env_var,
    get_genai_engine_version,
    get_logger,
    is_agentic_ui_enabled,
    is_api_only_mode_enabled,
    is_local_environment,
    is_transfer_encoding_middleware_enabled,
    new_relic_enabled,
    relevance_models_enabled,
)

logger = get_logger(log_level=Config.get_log_level())
logger.info(f"GenAI Engine log level set to: {logging._levelToName[logger.level]}")

tags_metadata = [
    {
        "name": "Rules",
        "description": "Endpoints to manage rules",
    },
    {
        "name": "Tasks",
        "description": "Endpoints to manage tasks and their rules",
    },
    {
        "name": "Default Validation",
        "description": "Endpoints to validate prompt and response on default rules",
    },
    {
        "name": "Task Based Validation",
        "description": "Endpoints to validate prompt and response for a task",
    },
    {
        "name": "Feedback",
        "description": "Endpoints to manage user feedback on inferences",
    },
    {
        "name": "Usage",
        "description": "Endpoints for retrieving token usage",
    },
    {
        "name": "Inferences",
        "description": "Endpoints for querying past inferences",
    },
    {
        "name": "API Keys",
        "description": "Endpoints for API keys management",
    },
    {
        "name": "Secrets",
        "description": "Endpoints for secrets management",
    },
    {
        "name": "Model Providers",
        "description": "Endpoints for model provider management",
    },
    {
        "name": "RAG Settings",
        "description": "Endpoints for RAG setting management",
    },
    {
        "name": "RAG Providers",
        "description": "Endpoints for RAG provider management",
    },
    {
        "name": "Prompt Experiments",
        "description": "Endpoints for managing prompt experiments",
    },
    {
        "name": "RAG Experiments",
        "description": "Endpoints for managing RAG experiments",
    },
    {
        "name": "Notebooks",
        "description": "Endpoints for managing experiment notebooks",
    },
]


def bootstrap_genai_engine_keycloak() -> None:
    lock_file = "/tmp/bootstrap.lock"
    with open(lock_file, "w") as f:
        # Acquire an exclusive lock (blocking)
        fcntl.flock(f, fcntl.LOCK_EX)
        try:
            logger.info("Bootstrapping GenAI Engine Keycloak")
            # Perform the bootstrap operation
            kc_client = get_keycloak_client().__next__()
            kc_client.bootstrap_genai_engine_keycloak()
        finally:
            logger.info("GenAI Engine Keycloak bootstrapped")
            # Release the lock
            fcntl.flock(f, fcntl.LOCK_UN)


def cleanup_cuda_cache() -> None:
    """Clean up PyTorch CUDA cache when the server is stopped."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        logger.info("Cleared PyTorch CUDA cache")


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    send_telemetry_event(TelemetryEventTypes.SERVER_START_INITIATED)

    device = torch.device(get_device())
    logger.info(f"Using device: {device.type}")

    try:
        db = next(get_db_session())

        # Initialize system task
        SystemTaskRepository(db).initialize_system_tasks()

        db.close()
    except HTTPException as e:
        raise ConnectionError(f"Error connecting to database: {e}") from None

    try:
        with db_session_context() as init_db:
            initialize_system_tasks(init_db)
    except Exception as e:
        logger.error(f"Error initializing system tasks: {e}")

    keycloak_settings = get_keycloak_settings()
    if keycloak_settings.ENABLED:
        bootstrap_genai_engine_keycloak()
    if not is_api_only_mode_enabled():
        get_oauth_client()
        time.sleep(random.randint(0, 3))

    # Download models in worker process
    logger.info("Downloading models...")
    try:
        model_load.download_models(1)  # Use single process in worker
    except Exception as e:
        logger.error(f"Error downloading models: {e}")
        raise e
    logger.info("Models downloaded.")

    model_load.get_claim_classifier_embedding_model()
    model_load.get_prompt_injection_model()
    model_load.get_prompt_injection_tokenizer()
    model_load.get_toxicity_model()
    model_load.get_toxicity_tokenizer()

    # Initialize continuous eval queue service
    try:
        initialize_continuous_eval_queue_service(num_workers=4)
    except Exception as e:
        logger.error(f"Error initializing continuous eval queue service: {e}")

    # Initialize currency conversion service (exchange rates, 6-hour refresh at 00/06/12/18 UTC)
    try:
        initialize_currency_conversion_service()
    except Exception as e:
        logger.error(f"Error initializing currency conversion service: {e}")

    # Initialize trace retention service (deletes trace data older than configured retention)
    try:
        initialize_trace_retention_service()
    except Exception as e:
        logger.error(f"Error initializing trace retention service: {e}")

    # Initialize global agent polling service
    try:
        initialize_global_agent_polling_service(num_workers=4)
    except Exception as e:
        logger.error(f"Error initializing global agent polling service: {e}")

    # Conditionally load relevance models
    if relevance_models_enabled():
        model_load.get_bert_scorer()
        model_load.get_relevance_reranker()
    else:
        logger.info(
            "Skipping relevance models loading - ENABLE_RELEVANCE_MODELS is False",
        )

    get_scorer_client()

    send_telemetry_event(TelemetryEventTypes.SERVER_START_COMPLETED)

    yield

    cleanup_cuda_cache()
    shutdown_trace_retention_service()
    shutdown_currency_conversion_service()
    shutdown_continuous_eval_queue_service()
    shutdown_global_agent_polling_service()


class TransferEncodingMiddleware(BaseHTTPMiddleware):
    async def dispatch(
        self,
        request: Request,
        call_next: RequestResponseEndpoint,
    ) -> Response:
        response = await call_next(request)
        # RFC 7230 §3.3.1 / RFC 9112 §6.1: must not send Transfer-Encoding
        # in 1xx or 204 responses
        if 100 <= response.status_code < 200 or response.status_code == 204:
            return response
        response.headers["Transfer-Encoding"] = "chunked"
        if "content-length" in response.headers:
            del response.headers["content-length"]
        return response


class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    def __get_default_csp_header_list(self, keycloak_uri: str) -> list[str]:
        headers = [
            # This is the fallback policy for everything else
            "default-src 'none'",
            "frame-ancestors 'none'",
            "base-uri 'self'",
            "form-action 'none'",
            # Whitelist things we need:
            # Fonts are served from our url /assets, so whitelist self for them
            "font-src 'self'",
            # This allows us to hit our API
            # Docs site uses some external images / scripts / inline scripts, manually whitelist them
            "img-src 'self' fastapi.tiangolo.com data:",
            "script-src 'self' *.jsdelivr.net 'sha256-QOOQu4W1oxGqd2nbXbxiA1Di6OHQOLQD+o+G9oWL8YY=' 'nonce-b6e054c5-d77e-4712-916e-1c027d548ac4' 'unsafe-inline'",
            # @todo make nonce hash generation as part of the build process
            "style-src 'self' 'nonce-b6e054c5-d77e-4712-916e-1c027d548ac4' cdn.jsdelivr.net",
            "report-uri /api/v2/csp_report",
            "trusted-types dompurify 'allow-duplicates'",
            # "require-trusted-types-for 'script'",
        ]
        if not is_local_environment():
            headers.append("upgrade-insecure-requests")
        if keycloak_uri:
            headers.append(f"connect-src 'self' {keycloak_uri}")
        else:
            headers.append("connect-src 'self'")
        return headers

    async def dispatch(
        self,
        request: Request,
        call_next: RequestResponseEndpoint,
    ) -> Response:
        response = await call_next(request)
        # Tell browsers to trust the content type header instead of trying to sniff the MIME type
        response.headers[constants.X_CONTENT_TYPE_OPTIONS_HEADER] = "nosniff"
        # Disallow browsers from loading GenAI Engine in an iframe
        response.headers[constants.X_FRAME_OPTIONS_HEADER] = "DENY"
        # Tell browsers to use only HTTPS to connect to this service for the next 2 years from request time
        response.headers[constants.STRICT_TRANSPORT_SECURITY_HEADER] = (
            "max-age=63072000; includeSubDomains"
        )

        keycloak_uri = ""

        # See this issue for why we can't apply the require-trusted-types-for 'script' policy to the entire site: https://github.com/acacode/swagger-typescript-api/issues/985
        response.headers[constants.CONTENT_SECURITY_POLICY_HEADER] = "; ".join(
            self.__get_default_csp_header_list(keycloak_uri),
        )
        response.headers[constants.CROSS_ORIGIN_OPENER_POLICY_HEADER] = "same-origin"
        response.headers[constants.CROSS_ORIGIN_EMBEDDER_POLICY_HEADER] = "require-corp"
        response.headers[constants.CROSS_ORIGIN_RESOURCE_POLICY_HEADER] = "same-origin"
        response.headers[constants.PERMISSIONS_POLICY_HEADER] = (
            "geolocation=(), camera=(), microphone=()"
        )
        response.headers[constants.REFERRER_POLICY_HEADER] = "no-referrer"
        return response


async def add_opentelemetry_custom_attributes(
    request: Request,
    call_next: RequestResponseEndpoint,
) -> Response:
    s = trace.get_current_span()
    if request.client is not None:
        s.set_attribute("http.client_ip", request.client.host)
    if "X-Forwarded-For" in request.headers:
        s.set_attribute("http.x_forwarded_for", request.headers["X-Forwarded-For"])
    response = await call_next(request)
    if s is not None:
        response.headers[constants.RESPONSE_TRACE_ID_HEADER] = hex(
            s.get_span_context().trace_id,
        ).removeprefix("0x")
    return response


def setup_newrelic(app: FastAPI) -> None:
    app.add_middleware(BaseHTTPMiddleware, dispatch=add_opentelemetry_custom_attributes)
    service_name = get_env_var(
        constants.NEWRELIC_APP_NAME_ENV_VAR,
        none_on_missing=False,
    )
    if service_name is None:
        service_name = "genai-engine"

    OTEL_RESOURCE_ATTRIBUTES: dict[str, str] = {"service.name": service_name}

    t = TracerProvider(resource=Resource.create(OTEL_RESOURCE_ATTRIBUTES))
    trace.set_tracer_provider(t)
    t.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))

    SQLAlchemyInstrumentor().instrument(tracer_provider=t, engine=get_db_engine())

    FastAPIInstrumentor.instrument_app(app)

    logger_provider = LoggerProvider(
        resource=Resource.create(attributes=OTEL_RESOURCE_ATTRIBUTES),
    )
    _logs.set_logger_provider(logger_provider)
    logger_provider.add_log_record_processor(
        BatchLogRecordProcessor(OTLPLogExporter()),
    )
    h = LoggingHandler(logger_provider=logger_provider)

    logger.addHandler(h)
    logging.getLogger("uvicorn.access").addHandler(h)


def get_base_app(
    version: str = get_genai_engine_version(),
    lifespan: Lifespan[FastAPI] | None = lifespan,
) -> FastAPI:
    logger.info(f"Booting up Arthur GenAI Engine: {version}")
    kwargs = {}
    if lifespan:
        kwargs["lifespan"] = lifespan
    app = FastAPI(
        title="Arthur GenAI Engine",
        version=version,
        openapi_tags=tags_metadata,
        **kwargs,  # type: ignore[arg-type]
    )
    origins = [
        "http://localhost",
        "http://0.0.0.0:8000",
        "http://0.0.0.0:3030",
        "http://localhost:3030",
        "http://localhost:8080",
        "http://localhost:3023",
        "http://localhost:3001",
        "http://localhost:3000",
    ]
    if ingress_url := get_env_var(
        constants.GENAI_ENGINE_INGRESS_URI_ENV_VAR,
        none_on_missing=True,
    ):
        origins.append(ingress_url)

    if cors_extra := get_env_var(
        constants.CORS_EXTRA_ORIGINS_ENV_VAR,
        none_on_missing=True,
    ):
        origins.extend(o.strip() for o in cors_extra.split(",") if o.strip())

    arthur_allowed_origins = r"https://.*\.arthur\.ai"

    app.add_middleware(
        CORSMiddleware,
        allow_origins=origins,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
        allow_origin_regex=arthur_allowed_origins,
    )
    return app


def add_routers(app: FastAPI, routers: list[APIRouter]) -> None:
    for router in routers:
        app.include_router(router)


def get_app_with_routes() -> FastAPI:
    app = get_base_app(lifespan=None)
    add_routers(
        app,
        [
            health_router,
            engine_config_routes,
            system_management_routes,
            feedback_routes,
            query_routes,
            rule_management_routes,
            task_management_routes,
            validate_routes,
            api_keys_routes,
            span_routes,
            trace_api_routes,
            agentic_prompt_routes,
            dataset_management_routes,
            model_provider_routes,
            secrets_routes,
            rag_routes,
            rag_setting_routes,
            llm_eval_routes,
            notebook_routes,
            rag_notebook_routes,
            agentic_notebook_routes,
            prompt_experiment_routes,
            rag_experiment_routes,
            agentic_experiment_routes,
            transform_routes,
            continuous_eval_routes,
            agent_polling_routes,
            demo_task_routes,
            tenant_signup_routes,
        ],
    )
    if extra_feature_config.CHATBOT_ENABLED:
        add_routers(app, [chatbot_routes])
    add_routers(app, [auth_routes, user_management_routes, user_identity_routes])
    add_routers(app, [app_chat_routes])
    return app


def get_test_app() -> FastAPI:
    app = get_base_app()
    # app.add_middleware(SecurityHeadersMiddleware)
    if is_transfer_encoding_middleware_enabled():
        app.add_middleware(TransferEncodingMiddleware)

    app.add_middleware(SessionMiddleware, secret_key=Config.app_secret_key())

    if Config.audit_log_enabled():
        setup_audit_logger()
        app.add_middleware(AuditLogMiddleware)

    add_routers(
        app,
        [
            health_router,
            engine_config_routes,
            system_management_routes,
            feedback_routes,
            query_routes,
            rule_management_routes,
            task_management_routes,
            validate_routes,
            api_keys_routes,
            span_routes,
            trace_api_routes,
            agentic_prompt_routes,
            dataset_management_routes,
            model_provider_routes,
            secrets_routes,
            rag_routes,
            rag_setting_routes,
            llm_eval_routes,
            notebook_routes,
            rag_notebook_routes,
            agentic_notebook_routes,
            prompt_experiment_routes,
            rag_experiment_routes,
            agentic_experiment_routes,
            transform_routes,
            continuous_eval_routes,
            agent_polling_routes,
            chatbot_routes,
            demo_task_routes,
            tenant_signup_routes,
        ],
    )
    add_routers(app, [auth_routes, user_management_routes, user_identity_routes])
    add_routers(app, [app_chat_routes])

    if is_api_only_mode_enabled():

        @app.get("/", include_in_schema=False)
        async def redirect_to_docs() -> RedirectResponse:
            return RedirectResponse("/docs")

    return app


def get_app() -> FastAPI:
    app = get_base_app()
    # app.add_middleware(SecurityHeadersMiddleware)
    if is_transfer_encoding_middleware_enabled():
        app.add_middleware(TransferEncodingMiddleware)

    app.add_middleware(SessionMiddleware, secret_key=Config.app_secret_key())

    if Config.audit_log_enabled():
        setup_audit_logger()
        app.add_middleware(AuditLogMiddleware)

    if new_relic_enabled():
        setup_newrelic(app)

    add_routers(
        app,
        [
            health_router,
            engine_config_routes,
            system_management_routes,
            feedback_routes,
            query_routes,
            rule_management_routes,
            task_management_routes,
            validate_routes,
            api_keys_routes,
            span_routes,
            trace_api_routes,
            agentic_prompt_routes,
            dataset_management_routes,
            model_provider_routes,
            secrets_routes,
            rag_routes,
            rag_setting_routes,
            llm_eval_routes,
            notebook_routes,
            rag_notebook_routes,
            agentic_notebook_routes,
            prompt_experiment_routes,
            rag_experiment_routes,
            agentic_experiment_routes,
            transform_routes,
            continuous_eval_routes,
            agent_polling_routes,
        ],
    )
    if extra_feature_config.CHATBOT_ENABLED:
        add_routers(app, [chatbot_routes])
    if Config.demo_mode():
        add_routers(app, [demo_task_routes])
    if extra_feature_config.CHAT_ENABLED:
        add_routers(app, [app_chat_routes])
    if extra_feature_config.DEMO_MODE:
        add_routers(app, [tenant_signup_routes])
    add_routers(app, [user_identity_routes])
    if not is_api_only_mode_enabled():
        add_routers(app, [auth_routes, user_management_routes])

    if is_agentic_ui_enabled():
        # Serve the React SPA
        static_dir = "/home/nonroot/app/static"
        if os.path.exists(static_dir):
            app.mount(
                "/",
                SPAStaticFiles(directory=static_dir, html=True),
                name="static",
            )

    return app


def start() -> None:
    send_telemetry_event(TelemetryEventTypes.SERVER_START_INITIATED)
    app = get_app()
    port = int(get_env_var(constants.GENAI_ENGINE_PORT_ENV_VAR, default="3030"))
    uvicorn.run(app, host="0.0.0.0", port=port)


if __name__ == "__main__":
    start()
