import json
from typing import Dict, Union
from mattermostdriver import Driver
from sqlalchemy import create_engine
from mnemosynecore.vault.client import (
    get_connection_as_json,
    get_connection_as_json_test,
)


def _build_connection(cfg: Dict, conn_type: str) -> Union[Driver, "Engine", Dict, object]:
    conn_type = conn_type.lower()
    login = cfg.get("user") or cfg.get("login")
    extra = cfg.get("extra", {})
    if isinstance(extra, str):
        try:
            extra = json.loads(extra or "{}")
        except Exception:
            extra = {}
    if not isinstance(extra, dict):
        extra = {}

    if conn_type in ("vertica", "vertica_engine"):
        if not login:
            raise ValueError("Для Vertica нужен login или user")
        vertica_url = (
            f"vertica+vertica_python://{login}:{cfg['password']}@"
            f"{cfg['host']}:{cfg['port']}/{cfg.get('schema')}"
        )
        return create_engine(vertica_url, pool_pre_ping=True)

    elif conn_type == "clickhouse":
        from clickhouse_driver import Client
        if not login:
            raise ValueError("Для ClickHouse нужен login или user")

        return Client(
            host=cfg["host"],
            port=int(cfg.get("port", 9000)),
            user=login,
            password=cfg.get("password"),
            database=cfg.get("schema") or cfg.get("database") or extra.get("database"),
            secure=cfg.get("secure", extra.get("secure", False)),
        )

    elif conn_type == "superset":
        return {
            "host": cfg["host"],
            "login": login,
            "password": cfg.get("password"),
            "schema": cfg.get("schema", "https"),
            "port": int(cfg.get("port", 443)),
            "extra": extra,
        }

    elif conn_type == "mattermost":
        driver = Driver({
            "url": cfg["host"],
            "token": cfg["password"],
            "scheme": cfg.get("schema", "https"),
            "port": int(cfg.get("port", 443)),
            "basepath": extra.get("basepath", "/api/v4"),
        })
        driver.login()
        return driver

    elif conn_type == "raw":
        return cfg

    else:
        raise ValueError(f"Неизвестный conn_type: {conn_type}")


def un_conn(conn_id: str, conn_type: str) -> Union[Driver, "Engine", Dict, object]:
    cfg = json.loads(get_connection_as_json(conn_id))
    return _build_connection(cfg, conn_type)


def un_conn_test(
    conn_id: str,
    conn_type: str,
    dir_path: str | None = None,
    creds_path: str | None = None,
) -> Union[Driver, "Engine", Dict, object]:
    cfg = json.loads(
        get_connection_as_json_test(
            conn_id=conn_id,
            dir_path=dir_path,
            creds_path=creds_path,
        )
    )
    return _build_connection(cfg, conn_type)
