import importlib
import json
import os
from pathlib import Path
import sys
import traceback


PROJECT_ROOT = Path(__file__).resolve().parents[2]
CACHE_ROOT = PROJECT_ROOT / ".cache"
cache_settings = {
    "LOCALAPPDATA": CACHE_ROOT / "localappdata",
    "MPLCONFIGDIR": CACHE_ROOT / "matplotlib",
    "NUMBA_CACHE_DIR": CACHE_ROOT / "numba",
    "SKB_DATA_DIRECTORY": CACHE_ROOT / "skrub",
    "HF_HOME": CACHE_ROOT / "huggingface",
    "TORCH_HOME": CACHE_ROOT / "torch",
    "XDG_CACHE_HOME": CACHE_ROOT,
}
for key, path in cache_settings.items():
    path.mkdir(parents=True, exist_ok=True)
    os.environ[key] = str(path)
os.environ["PYTENSOR_FLAGS"] = f"base_compiledir={CACHE_ROOT / 'pytensor'}"

import numpy as np


results = {"python": sys.version, "imports": {}, "tests": {}}


def record(name, fn):
    try:
        value = fn()
        results["tests"][name] = {"status": "passed", "detail": value}
    except Exception as exc:
        results["tests"][name] = {
            "status": "failed",
            "detail": f"{type(exc).__name__}: {exc}",
            "traceback": traceback.format_exc(),
        }


for module_name in [
    "numpy", "scipy", "pandas", "polars", "pyarrow", "duckdb",
    "statsmodels", "pymc", "pytensor", "arviz", "nutpie",
    "jax", "numpyro", "blackjax", "bambi", "semopy",
    "factor_analyzer", "pingouin", "girth", "torch", "torchvision",
    "transformers", "sentence_transformers", "xgboost",
    "tabpfn",
]:
    try:
        module = importlib.import_module(module_name)
        results["imports"][module_name] = {
            "status": "passed",
            "version": getattr(module, "__version__", "unknown"),
        }
    except Exception as exc:
        results["imports"][module_name] = {
            "status": "failed",
            "detail": f"{type(exc).__name__}: {exc}",
        }


def pymc_test():
    import pymc as pm

    rng = np.random.default_rng(20260920)
    observed = rng.normal(0.35, 1.0, size=40)
    with pm.Model():
        mu = pm.Normal("mu", 0.0, 1.0)
        pm.Normal("y", mu, 1.0, observed=observed)
        inference = pm.sample(
            draws=50,
            tune=50,
            chains=1,
            cores=1,
            random_seed=20260920,
            progressbar=False,
            compute_convergence_checks=False,
        )
    posterior_mean = float(inference.posterior["mu"].mean())
    return {"posterior_mean": posterior_mean, "draws": 50, "chains": 1}


def sem_test():
    import pandas as pd
    import semopy

    rng = np.random.default_rng(20260920)
    latent = rng.normal(size=160)
    frame = pd.DataFrame(
        {
            "x1": 0.8 * latent + rng.normal(scale=0.4, size=160),
            "x2": 0.7 * latent + rng.normal(scale=0.5, size=160),
            "x3": 0.9 * latent + rng.normal(scale=0.3, size=160),
        }
    )
    model = semopy.Model("factor =~ x1 + x2 + x3")
    fit = model.fit(frame)
    return {"success": bool(fit.success), "objective": float(fit.fun)}


def factor_test():
    from factor_analyzer import FactorAnalyzer

    rng = np.random.default_rng(20260920)
    latent = rng.normal(size=(180, 1))
    data = latent @ np.array([[0.8, 0.7, 0.9, 0.6]]) + rng.normal(
        scale=0.45, size=(180, 4)
    )
    fitted = FactorAnalyzer(n_factors=1, rotation=None).fit(data)
    return {"loading_shape": list(fitted.loadings_.shape)}


def torch_test():
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("torch.cuda.is_available() is false")
    left = torch.randn((256, 256), device="cuda")
    right = torch.randn((256, 256), device="cuda")
    result = left @ right
    torch.cuda.synchronize()
    return {
        "torch": torch.__version__,
        "cuda_runtime": torch.version.cuda,
        "device": torch.cuda.get_device_name(0),
        "finite": bool(torch.isfinite(result).all().item()),
    }


record("pymc_sampling", pymc_test)
record("semopy_cfa", sem_test)
record("factor_analyzer", factor_test)
record("torch_cuda", torch_test)

serialized = json.dumps(results, ensure_ascii=False, indent=2)
(Path(__file__).parent / "smoke-test-result.json").write_text(
    serialized, encoding="utf-8"
)
print(serialized)

failed_imports = [k for k, v in results["imports"].items() if v["status"] != "passed"]
failed_tests = [k for k, v in results["tests"].items() if v["status"] != "passed"]
raise SystemExit(1 if failed_imports or failed_tests else 0)