#!/usr/bin/env python3

# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

import os
import sys
import runpy
import importlib
import subprocess
from pathlib import Path

from datatailr.logging import DatatailrLogger
from datatailr.scheduler.utils import get_imports_from_code


logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()

SUPPORTED_FRAMEWORKS = {"streamlit", "flask", "fastapi", "dash", "panel"}


def _parse_entrypoint(entrypoint: str) -> tuple[str, str]:
    if ":" not in entrypoint:
        raise ValueError(
            "Environment variable 'DATATAILR_ENTRYPOINT' must be in the format 'module:attr'."
        )
    module_name, attr_name = entrypoint.split(":", 1)
    return module_name, attr_name


def _detect_framework_from_imports(script_path: str) -> str | None:
    imports, top_level_packages = get_imports_from_code(
        script_path, include_relatives=True
    )

    if "streamlit" in top_level_packages:
        return "streamlit"
    if "flask" in top_level_packages:
        return "flask"
    if "fastapi" in top_level_packages:
        return "fastapi"
    if "dash" in top_level_packages:
        return "dash"
    if "panel" in top_level_packages:
        return "panel"

    return None


def run() -> None:
    logger.info("Starting Datatailr app...")

    entrypoint = os.environ.get("DATATAILR_ENTRYPOINT")
    if not entrypoint:
        raise ValueError(
            "Environment variable 'DATATAILR_ENTRYPOINT' must be set (module:attr)."
        )

    module_name, attr_name = _parse_entrypoint(entrypoint)

    # Host/port/worker configuration with sensible defaults
    host = "0.0.0.0"
    port = "8080"
    workers = "1"  # only relevant for flask using gunicorn.

    # Explicit framework selection takes precedence
    framework = os.environ.get("DATATAILR_FRAMEWORK")

    spec = importlib.util.find_spec(module_name)
    if spec is None:
        raise ImportError(f"Module '{module_name}' not found for entrypoint.")
    if spec.origin is None:
        raise ImportError(f"Module '{module_name}' does not have a file origin.")
    script_path = Path(spec.origin)
    if not framework:
        framework = _detect_framework_from_imports(str(script_path))

    if framework not in SUPPORTED_FRAMEWORKS:
        raise ValueError(
            f"Unsupported or undetected framework '{framework}'. "
            "Set DATATAILR_FRAMEWORK to one of: "
            f"{', '.join(sorted(SUPPORTED_FRAMEWORKS))}."
        )

    logger.info(
        "Running entrypoint '%s' as framework '%s' on %s:%s",
        entrypoint,
        framework,
        host,
        port,
    )

    if framework == "streamlit":
        sys.argv = [
            "streamlit",
            "run",
            str(script_path),
            "--server.address",
            host,
            "--server.port",
            port,
            "--server.headless",
            "true",
            *sys.argv[1:],
        ]
        runpy.run_module("streamlit", run_name="__main__")

    elif framework == "panel":
        # Find the actual app file in the module. One of 'app.py' or 'main.py' should exist.
        app_files = ["app.py", "main.py"]
        for app_file in app_files:
            if os.path.exists(os.path.join(os.path.dirname(script_path), app_file)):
                app_file = os.path.join(os.path.dirname(script_path), app_file)
                break
        else:
            raise FileNotFoundError(
                f"App file '{app_files}' not found in module '{module_name}'."
            )

        logger.info("Starting Panel app: %s", app_file)
        sys.argv = [
            "panel",
            "serve",
            str(app_file),
            "--allow-websocket-origin",
            "*",
            "--address",
            host,
            "--port",
            port,
            *sys.argv[1:],
        ]
        runpy.run_module("panel", run_name="__main__")

    elif framework == "flask":
        cmd = [
            "gunicorn",
            "--bind",
            f"{host}:{port}",
            f"{module_name}:app",
        ]
        if workers:
            cmd.extend(["--workers", workers])

        logger.info("Starting gunicorn WSGI server: %s", " ".join(cmd))
        subprocess.run(cmd, check=True)

    elif framework == "dash":
        # Set DASH_REQUESTS_PATHNAME_PREFIX before importing the user module so
        # the Dash app picks it up at construction time.
        # This has to be done before importing the user module so that the Dash app
        # picks it up at construction time.
        # The prefix is: /job/{env}/{job_name}/ where env and name can be obtained from environment variables.
        env = os.environ.get("DATATAILR_JOB_ENVIRONMENT", "dev")
        job_name = os.environ.get("DATATAILR_JOB_NAME", "default")
        os.environ["DASH_REQUESTS_PATHNAME_PREFIX"] = f"/job/{env}/{job_name}/"

        spec = importlib.util.find_spec(module_name)
        if spec is None:
            raise ImportError(f"Module '{module_name}' not found for entrypoint.")
        module = importlib.util.module_from_spec(spec)
        if spec.loader is None:
            raise ImportError(f"Module '{module_name}' does not have a loader.")
        spec.loader.exec_module(module)
        server = getattr(module, "server", None)
        if server is None:
            raise ValueError(
                "Dash entrypoint must expose a Dash app with a 'server' attribute for WSGI."
            )

        wsgi_app = f"{module_name}:server"

        cmd = [
            "gunicorn",
            "--bind",
            f"{host}:{port}",
            wsgi_app,
        ]
        if workers:
            cmd.extend(["--workers", workers])

        logger.info("Starting gunicorn WSGI server for Dash: %s", " ".join(cmd))
        subprocess.run(cmd, check=True)

    elif framework == "fastapi":
        env = os.environ.get("DATATAILR_JOB_ENVIRONMENT", "dev")
        job_name = os.environ.get("DATATAILR_JOB_NAME", "default")
        root_path = f"/job/{env}/{job_name}"

        cmd = [
            "uvicorn",
            f"{module_name}:app",
            "--host",
            host,
            "--port",
            port,
            "--root-path",
            root_path,
            "--proxy-headers",
            "--forwarded-allow-ips",
            "*",
        ]
        logger.info("Starting uvicorn ASGI server: %s", " ".join(cmd))
        subprocess.run(cmd, check=True)

    else:
        raise RuntimeError(f"Unhandled framework branch: {framework}")
