"""Access wordslab-notebooks environment version, platform, urls, ports, directories (install / apps / data / models), and pre-installed models."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_env.ipynb.

# %% auto 0
__all__ = ['WordslabEnv']

# %% ../nbs/01_env.ipynb 2
import os, re, subprocess
from pathlib import Path
from typing import Optional

from fastcore.utils import patch

# %% ../nbs/01_env.ipynb 5
class WordslabEnv:
    def __init__(self):
        # wordslab-notebooks version and platform
        self.version = os.environ["WORDSLAB_VERSION"]
        self.platform = os.environ["WORDSLAB_PLATFORM"]

        # wordslab-notebooks external urls
        self.url_dashboard = os.environ["DASHBOARD_URL"]
        self.url_openwebui = os.environ["OPENWEBUI_URL"]
        self.url_jupyterlab = os.environ["JUPYTERLAB_URL"]
        self.url_vscode = os.environ["VSCODE_URL"]
        self.url_userapp1 = os.environ["USER_APP1_URL"]
        self.url_userapp2 = os.environ["USER_APP2_URL"]
        self.url_userapp3 = os.environ["USER_APP3_URL"]
        self.url_userapp4 = os.environ["USER_APP4_URL"]
        self.url_userapp5 = os.environ["USER_APP5_URL"]

        # wordslab-notebooks internal ports
        self.port_ollama = 11434
        self.port_vllm = 8000
        self.port_docling = 5001   
        self.port_userapp1 = os.environ["USER_APP1_PORT"]
        self.port_userapp2 = os.environ["USER_APP2_PORT"]
        self.port_userapp3 = os.environ["USER_APP3_PORT"]
        self.port_userapp4 = os.environ["USER_APP4_PORT"]
        self.port_userapp5 = os.environ["USER_APP5_PORT"]
        
        # wordslab-notebooks install directories
        self.dir_home = os.environ["WORDSLAB_HOME"]
        self.dir_scripts = os.environ["WORDSLAB_SCRIPTS"]     
        self.dir_python = os.environ["UV_PYTHON_INSTALL_DIR"]  
        self.dir_workspace = os.environ["WORDSLAB_WORKSPACE"]
        self.dir_models = os.environ["WORDSLAB_MODELS"]

        # wordslab-notebooks applications install directories
        self.dir_openwebui = os.environ["OPENWEBUI_ENV"]
        self.dir_jupyterlab = os.environ["JUPYTERLAB_ENV"]
        self.dir_vscode = os.environ["VSCODE_DIR"]
        self.dir_ollama = os.environ["OLLAMA_DIR"]
        self.dir_docling = os.environ["DOCLING_ENV"]  

        # wordslab-notebooks applications data directories
        self.dir_openwebui_data = os.environ["OPENWEBUI_DATA"]
        self.dir_jupyterlab_data = os.environ["JUPYTER_DATA"]
        self.dir_vscode_data = os.environ["VSCODE_DATA"]
        self.dir_docling_data = os.environ["DOCLING_DATA"]
        
        # wordslab-notebooks models directories
        self.dir_models_ollama = os.environ["OLLAMA_MODELS"] 
        self.dir_models_vllm = os.environ["HF_HOME"]
        self.dir_models_hugginface = os.environ["HF_HOME"]
        self.dir_models_fastai = os.environ["FASTAI_HOME"] 
        self.dir_models_pytorch = os.environ["TORCH_HOME"] 
        self.dir_models_keras = os.environ["KERAS_HOME"]
        self.dir_models_tensorflow = os.environ["TFHUB_CACHE_DIR"]         
        self.dir_models_docling = os.environ["DOCLING_MODELS"]
        
        # wordslab-notebooks default models
        self.default_model_chat = os.environ["OLLAMA_CHAT_MODEL"]
        self.default_model_code = os.environ["OLLAMA_CODE_MODEL"]
        self.default_model_agent = os.environ["OLLAMA_AGENT_MODEL"]        
        self.default_model_embedding = os.environ["OLLAMA_EMBED_MODEL"]
        self.default_model_context_length = os.environ["OLLAMA_CONTEXT_LENGTH"]

        # external cloud services - api keys
        # Note: you need to reload the WordslabEnv() object after you define one of these environment variables
        self.cloud_openrouter_api_key_var = "OPENROUTER_API_KEY"
        self.cloud_openrouter_api_key = os.environ.get(self.cloud_openrouter_api_key_var)
        self.cloud_replicate_api_token_var = "REPLICATE_API_TOKEN"
        self.cloud_replicate_api_token = os.environ.get(self.cloud_replicate_api_token_var)
        self.cloud_tavily_api_key_var = "TAVILY_API_KEY"
        self.cloud_tavily_api_key = os.environ.get(self.cloud_tavily_api_key_var)
        self.cloud_ollama_api_key_var = "OLLAMA_API_KEY"
        self.cloud_ollama_api_key = os.environ.get(self.cloud_ollama_api_key_var)
        self.cloud_huggingface_access_token_var = "HF_TOKEN"
        self.cloud_huggingface_access_token = os.environ.get(self.cloud_huggingface_access_token_var)

# %% ../nbs/01_env.ipynb 63
_ENV_REF_PATTERN = re.compile(
    r"""
    \$(\w+)|           # $VAR
    \${([^}]+)}        # ${VAR}
    """,
    re.VERBOSE,
)

def _expand_with_current_env(value: str) -> str:
    """
    Expand $VAR and ${VAR} using os.environ.
    Unknown variables are left unchanged.
    """

    def replacer(match: re.Match) -> str:
        var_name = match.group(1) or match.group(2)
        return os.environ.get(var_name, match.group(0))

    return _ENV_REF_PATTERN.sub(replacer, value)

@patch
def write_user_env_var(self: WordslabEnv, var_name: str, var_value: str) -> None:
    """
    Write or update an environment variable in
    $WORDSLAB_WORKSPACE/.secrets/user-env.bashrc

    The file is created if it doesn't exist.
    The variable is created if it doesn't exist and updated if it already exist.
    The value is safely quoted for Bash while preserving $VAR references.

    The new variable value is automatically injected in the current python process.
    The environment variable will be automatically loaded each time you start a new shell.
    But you need to restart all the currently running shells and other processes to see the new variable or value.
    """

    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", var_name):
        raise ValueError(f"Invalid environment variable name: {var_name}")

    workspace = os.environ.get("WORDSLAB_WORKSPACE")
    if not workspace:
        raise EnvironmentError("WORDSLAB_WORKSPACE is not set")

    env_file = Path(workspace) / ".secrets" / "user-env.bashrc"
    env_file.parent.mkdir(parents=True, exist_ok=True)

    # Escape value for double-quoted Bash string (but keep $ unescaped)
    escaped_value = (
        var_value
        .replace("\\", "\\\\")
        .replace('"', '\\"')
        .replace("`", "\\`")
    )

    new_line = f'export {var_name}="{escaped_value}"'

    if env_file.exists():
        lines = env_file.read_text().splitlines()
    else:
        lines = []

    var_pattern = re.compile(rf"^\s*(?:export\s+)?{re.escape(var_name)}\s*=")
    updated = False

    for i, line in enumerate(lines):
        if var_pattern.match(line):
            lines[i] = new_line
            updated = True
            break

    if not updated:
        if lines and lines[-1].strip() != "":
            lines.append("")
        lines.append(new_line)

    env_file.write_text("\n".join(lines) + "\n")

    expanded_value = _expand_with_current_env(var_value)
    os.environ[var_name] = expanded_value

# %% ../nbs/01_env.ipynb 64
@patch
def read_user_env_var(self: WordslabEnv, var_name: str) -> Optional[str]:
    """
    Read a variable value from
    $WORDSLAB_WORKSPACE/.secrets/user-env.bashrc

    Returns the raw (unexpanded) value, or None if not found.
    """

    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", var_name):
        raise ValueError(f"Invalid environment variable name: {var_name}")

    workspace = os.environ.get("WORDSLAB_WORKSPACE")
    if not workspace:
        raise EnvironmentError("WORDSLAB_WORKSPACE is not set")

    env_file = Path(workspace) / ".secrets" / "user-env.bashrc"
    if not env_file.exists():
        return None

    pattern = re.compile(
        rf"""
        ^\s*
        (?:export\s+)?
        {re.escape(var_name)}
        \s*=\s*
        (?:
            "(?P<dq>(?:\\.|[^"])*)" |
            (?P<bare>[^\s#]+)
        )
        """,
        re.VERBOSE,
    )

    for line in env_file.read_text().splitlines():
        if line.strip().startswith("#"):
            continue

        match = pattern.match(line)
        if match:
            if match.group("dq") is not None:
                return bytes(match.group("dq"), "utf-8").decode("unicode_escape")
            return match.group("bare")

    return None

# %% ../nbs/01_env.ipynb 71
@patch
def setup_openrouter(self: WordslabEnv, openrouter_api_key: str) -> None:
    self.write_user_env_var(self.cloud_openrouter_api_key_var, openrouter_api_key)

# %% ../nbs/01_env.ipynb 75
@patch
def setup_replicate(self: WordslabEnv, replicate_api_token: str) -> None:
    self.write_user_env_var(self.cloud_replicate_api_token_var, replicate_api_token)

# %% ../nbs/01_env.ipynb 79
@patch
def setup_tavily(self: WordslabEnv, tavily_api_key: str) -> None:
    self.write_user_env_var(self.cloud_tavily_api_key_var, tavily_api_key)

# %% ../nbs/01_env.ipynb 83
@patch
def setup_ollamacloud(self: WordslabEnv, ollama_api_key: str) -> None:
    self.write_user_env_var(self.cloud_ollama_api_key_var, ollama_api_key)

# %% ../nbs/01_env.ipynb 87
@patch
def setup_github(self: WordslabEnv, git_user_name: str, git_user_email: str, github_username: str, github_access_token: str) -> None:
    # Set global Git user name and email
    subprocess.run(["git", "config", "--global", "user.name", git_user_name], check=True)
    subprocess.run(["git", "config", "--global", "user.email", git_user_email], check=True)
    # Enable Git credential store
    subprocess.run(["git", "config", "--global", "credential.helper", "store"], check=True)
    # Store credentials in ~/.git-credentials
    credentials_file = Path.home() / ".git-credentials"
    credential_line = f"https://{github_username}:{github_access_token}@github.com\n"
    with open(credentials_file, "w") as f:
        f.write(credential_line)
    print("Wrote Github credentials");

# %% ../nbs/01_env.ipynb 90
@patch
def setup_huggingface(self: WordslabEnv, huggingface_access_token: str) -> None:
    self.write_user_env_var(self.cloud_huggingface_access_token_var, huggingface_access_token)

# %% ../nbs/01_env.ipynb 94
@patch
def setup_pypi(self: WordslabEnv, pypi_api_token: str) -> None:
    pypirc_path = Path.home() / ".pypirc"
    content = f"""[pypi]
username = __token__
password = {pypi_api_token}
"""
    pypirc_path.write_text(content)
    print("Wrote Pypi credentials");
