"""MLflow deployment target - log model + create Databricks serving endpoint."""

from __future__ import annotations

import getpass
import importlib.resources
import json
import tempfile
from collections.abc import Callable
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING

from haute.deploy._config import ResolvedDeploy

# Resolve the path to the models-from-code script shipped with the package
_MODEL_CODE_PATH = str(importlib.resources.files("haute.deploy") / "_model_code.py")

if TYPE_CHECKING:
    from haute.deploy._config import DeployConfig


@dataclass
class DeployResult:
    """Returned by deploy_to_mlflow()."""

    model_name: str
    model_version: int
    model_uri: str
    endpoint_url: str | None
    manifest_path: Path


def deploy_to_mlflow(
    resolved: ResolvedDeploy,
    progress: Callable[[str], None] | None = None,
) -> DeployResult:
    """Deploy a resolved pipeline to MLflow + Databricks Model Serving.

    Steps:
        1. Build deployment manifest JSON
        2. Log HauteModel as mlflow.pyfunc with artifacts + signature
        3. Register model version in MLflow Model Registry
        4. Return DeployResult with model URI and endpoint URL

    Args:
        resolved: Fully resolved deployment config (from ``resolve_config()``).
        progress: Optional callback for step-by-step progress messages.

    Returns:
        DeployResult with model URI, version, and endpoint URL.
    """

    def _log(msg: str) -> None:
        if progress:
            progress(msg)

    import mlflow

    config = resolved.config
    model_name = config.model_name

    # Point MLflow at the Databricks workspace (uses DATABRICKS_HOST/TOKEN env vars)
    _log("Connecting to Databricks MLflow...")
    _check_databricks_connectivity(_log)
    mlflow.set_tracking_uri("databricks")
    mlflow.set_registry_uri("databricks-uc")

    # Use Unity Catalog three-level namespace: catalog.schema.model_name
    uc_model_name = f"{config.databricks.catalog}.{config.databricks.schema}.{model_name}"

    # 1. Build deployment manifest
    manifest = _build_manifest(resolved)

    # 2. Write manifest to a temp file
    with tempfile.TemporaryDirectory() as tmpdir:
        manifest_path = Path(tmpdir) / "deploy_manifest.json"
        manifest_path.write_text(json.dumps(manifest, indent=2))

        # 3. Build artifact dict for mlflow.pyfunc.log_model
        artifacts: dict[str, str] = {
            "deploy_manifest": str(manifest_path),
        }
        for artifact_name, artifact_path in resolved.artifacts.items():
            artifacts[artifact_name] = str(artifact_path)

        # 4. Build MLflow model signature
        signature = _build_signature(resolved)

        # 5. Set experiment - append endpoint suffix for staging isolation
        experiment_name = config.databricks.experiment_name
        if config.endpoint_suffix:
            experiment_name = experiment_name + config.endpoint_suffix
        _log(f"Setting experiment: {experiment_name}")
        mlflow.set_experiment(experiment_name)

        # 6. Log the model
        _log("Logging model to MLflow (this may take a minute)...")
        with mlflow.start_run(run_name=f"deploy-{model_name}"):
            mlflow.log_dict(manifest, "deploy_manifest.json")

            mlflow.pyfunc.log_model(
                name="model",
                python_model=_MODEL_CODE_PATH,
                artifacts=artifacts,
                signature=signature,
                pip_requirements=_pip_requirements(resolved),
                registered_model_name=uc_model_name,
            )

        _log(f"Model logged. Fetching registered version for {uc_model_name}...")
        # 7. Get the registered model version
        client = mlflow.tracking.MlflowClient()
        versions = client.search_model_versions(f"name='{uc_model_name}'")
        if versions:
            latest_version = max(v.version for v in versions)
        else:
            latest_version = 1

        model_uri = f"models:/{uc_model_name}/{latest_version}"

    _log(f"Model URI: {model_uri}")

    # 8. Create or update the serving endpoint
    _log(f"Creating/updating serving endpoint: {config.effective_endpoint_name}...")
    endpoint_url = _create_or_update_serving_endpoint(
        config=config,
        uc_model_name=uc_model_name,
        model_version=int(latest_version),
    )

    return DeployResult(
        model_name=model_name,
        model_version=int(latest_version),
        model_uri=model_uri,
        endpoint_url=endpoint_url,
        manifest_path=manifest_path,
    )


def get_deploy_status(
    model_name: str,
    catalog: str = "main",
    schema: str = "default",
) -> dict[str, str | int]:
    """Query MLflow Model Registry for current model versions and status.

    Args:
        model_name: Short model name (e.g. ``"motor-pricing"``).
        catalog: Unity Catalog catalog name.
        schema: Unity Catalog schema name.

    Returns:
        Dict with keys: model_name, latest_version, latest_stage, status.
    """
    import mlflow

    mlflow.set_tracking_uri("databricks")
    mlflow.set_registry_uri("databricks-uc")

    uc_model_name = f"{catalog}.{schema}.{model_name}"
    client = mlflow.tracking.MlflowClient()
    versions = client.search_model_versions(f"name='{uc_model_name}'")

    if not versions:
        return {
            "model_name": model_name,
            "latest_version": 0,
            "status": "not_found",
        }

    latest = max(versions, key=lambda v: int(v.version))
    return {
        "model_name": model_name,
        "latest_version": int(latest.version),
        "latest_stage": getattr(latest, "current_stage", "None"),
        "status": latest.status,
        "run_id": latest.run_id,
    }


def _build_manifest(resolved: ResolvedDeploy) -> dict:
    """Build the deployment manifest dict."""
    import haute

    config = resolved.config
    return {
        "haute_version": haute.__version__,
        "pipeline_name": resolved.pruned_graph.get("pipeline_name", config.model_name),
        "pipeline_file": str(config.pipeline_file),
        "created_at": datetime.now(UTC).isoformat(),
        "created_by": _get_user(),
        "input_nodes": resolved.input_node_ids,
        "output_node": resolved.output_node_id,
        "output_fields": config.output_fields,
        "input_schema": resolved.input_schema,
        "output_schema": resolved.output_schema,
        "artifacts": {name: str(path) for name, path in resolved.artifacts.items()},
        "graph": resolved.pruned_graph,
        "nodes_deployed": len(resolved.pruned_graph.get("nodes", [])),
        "nodes_skipped": len(resolved.removed_node_ids),
        "nodes_skipped_names": resolved.removed_node_ids,
    }


def _build_signature(resolved: ResolvedDeploy) -> object:
    """Build an MLflow ModelSignature from resolved schemas."""
    from mlflow.models import ModelSignature
    from mlflow.types import ColSpec, DataType, Schema

    dtype_map = {
        "Int8": DataType.integer,
        "Int16": DataType.integer,
        "Int32": DataType.integer,
        "Int64": DataType.long,
        "UInt8": DataType.integer,
        "UInt16": DataType.integer,
        "UInt32": DataType.long,
        "UInt64": DataType.long,
        "Float32": DataType.float,
        "Float64": DataType.double,
        "String": DataType.string,
        "Utf8": DataType.string,
        "Boolean": DataType.boolean,
        "Date": DataType.datetime,
        "Datetime": DataType.datetime,
    }

    def _to_colspecs(schema: dict[str, str]) -> list[ColSpec]:
        specs = []
        for col_name, dtype_str in schema.items():
            # Handle parameterized types like Datetime('us', 'UTC')
            base_type = dtype_str.split("(")[0] if "(" in dtype_str else dtype_str
            mlflow_type = dtype_map.get(base_type, DataType.string)
            specs.append(ColSpec(type=mlflow_type, name=col_name))
        return specs

    input_schema = Schema(_to_colspecs(resolved.input_schema))
    output_schema = Schema(_to_colspecs(resolved.output_schema))
    return ModelSignature(inputs=input_schema, outputs=output_schema)


def _pip_requirements(resolved: ResolvedDeploy) -> list[str]:
    """Build pip requirements for the deployed model."""
    import haute

    reqs = [
        f"haute=={haute.__version__}",
        "polars>=1.0.0",
    ]

    # Check if catboost is used
    for node in resolved.pruned_graph.get("nodes", []):
        config = node.get("data", {}).get("config", {})
        if config.get("fileType") == "catboost":
            reqs.append("catboost>=1.2.8")
            break

    return reqs


def _create_or_update_serving_endpoint(
    config: DeployConfig,
    uc_model_name: str,
    model_version: int,
) -> str | None:
    """Create or update a Databricks Model Serving endpoint.

    Uses the Databricks SDK to create an endpoint if it doesn't exist,
    or update the served model version if it does.

    Args:
        config: The deployment configuration (provides endpoint_name and
                serving settings from haute.toml).
        uc_model_name: The Unity Catalog three-level model name
                       (e.g. ``"workspace.default.haute-test"``).
        model_version: The model version number to serve.

    Returns:
        The endpoint invocation URL, or ``None`` if endpoint_name is not
        configured.
    """
    import os

    endpoint_name = config.effective_endpoint_name
    if not endpoint_name:
        return None

    from databricks.sdk import WorkspaceClient
    from databricks.sdk.errors import NotFound
    from databricks.sdk.service.serving import (
        EndpointCoreConfigInput,
        ServedEntityInput,
    )

    host = os.environ.get("DATABRICKS_HOST", "").rstrip("/")

    ws = WorkspaceClient()

    served_entity = ServedEntityInput(
        entity_name=uc_model_name,
        entity_version=str(model_version),
        workload_size=config.databricks.serving_workload_size,
        scale_to_zero_enabled=config.databricks.serving_scale_to_zero,
    )

    try:
        ws.serving_endpoints.get(endpoint_name)
        # Endpoint exists - update the served model version
        ws.serving_endpoints.update_config(
            name=endpoint_name,
            served_entities=[served_entity],
        )
    except NotFound:
        # Endpoint doesn't exist - create it
        ws.serving_endpoints.create(
            name=endpoint_name,
            config=EndpointCoreConfigInput(
                name=endpoint_name,
                served_entities=[served_entity],
            ),
        )

    return f"{host}/serving-endpoints/{endpoint_name}/invocations"


def _check_databricks_connectivity(
    _log: Callable[[str], None],
    timeout: int = 10,
) -> None:
    """Verify the Databricks workspace is reachable before slow MLflow calls.

    Makes a lightweight HTTP GET to DATABRICKS_HOST with a short timeout so CI
    fails fast with a clear error instead of hanging forever.
    """
    import os
    import urllib.request

    host = os.environ.get("DATABRICKS_HOST", "")
    token = os.environ.get("DATABRICKS_TOKEN", "")

    if not host:
        raise RuntimeError("DATABRICKS_HOST is not set. Add it to .env or set it as a CI secret.")
    if not token:
        raise RuntimeError("DATABRICKS_TOKEN is not set. Add it to .env or set it as a CI secret.")

    url = f"{host.rstrip('/')}/api/2.0/clusters/list-zones"
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"})
    try:
        urllib.request.urlopen(req, timeout=timeout)  # noqa: S310
        _log("Databricks workspace reachable")
    except urllib.error.HTTPError as exc:
        if exc.code == 403:
            raise RuntimeError(
                f"Databricks returned 403 Forbidden. "
                f"Check that your DATABRICKS_TOKEN is valid and has workspace access. "
                f"Host: {host}"
            ) from exc
        # Other HTTP errors (e.g. 404) are fine — it means the host is reachable
        _log("Databricks workspace reachable")
    except (urllib.error.URLError, TimeoutError) as exc:
        raise RuntimeError(
            f"Cannot reach Databricks workspace at {host} "
            f"(timed out after {timeout}s). Check DATABRICKS_HOST is correct "
            f"and that the workspace allows connections from this network "
            f"(e.g. GitHub Actions IP ranges may need allowlisting)."
        ) from exc


def _get_user() -> str:
    """Get the current user's name."""
    try:
        return getpass.getuser()
    except (KeyError, OSError):
        return "unknown"
