import json
import logging
import shlex
import ssl
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional

import websockets
import yaml


log = logging.getLogger("truenas_middleware")


@dataclass
class TrueNASConfig:
    ws_url: str
    api_key: str
    api_user: Optional[str]
    app_name: str
    verify_ssl: bool = False


def _parse_compose(raw: Any) -> Dict[str, Any]:
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, str):
        text = raw.strip()
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return yaml.safe_load(text)
    raise ValueError("Unsupported compose payload")


def _command_to_list(command: Any) -> list:
    if isinstance(command, list):
        return command
    if isinstance(command, str):
        return shlex.split(command)
    return []


def _extract_command(config: Dict[str, Any], service_name: str = "llamacpp") -> list:
    if config.get("custom_compose_config") or config.get("custom_compose_config_string"):
        compose = _parse_compose(config.get("custom_compose_config") or config.get("custom_compose_config_string") or {})
        services = compose.get("services") or {}
        svc = services.get(service_name) or {}
        return _command_to_list(svc.get("command"))
    return _command_to_list(config.get("command"))


def _model_id_from_command(cmd: list) -> Optional[str]:
    if "--model" in cmd:
        idx = cmd.index("--model")
        if idx + 1 < len(cmd):
            return Path(cmd[idx + 1]).name
    return None


def _set_arg(cmd: list, flag: str, value: Optional[str]) -> list:
    if value is None:
        return cmd
    if flag in cmd:
        idx = cmd.index(flag)
        if idx + 1 < len(cmd):
            cmd[idx + 1] = value
        else:
            cmd.append(value)
        return cmd
    cmd.extend([flag, value])
    return cmd


def _merge_args(cmd: list, args: Dict[str, str]) -> list:
    flag_map = {
        "device": "--device",
        "tensor_split": "--tensor-split",
        "split_mode": "--split-mode",
        "n_gpu_layers": "--n-gpu-layers",
        "ctx_size": "--ctx-size",
        "batch_size": "--batch-size",
        "ubatch_size": "--ubatch-size",
        "cache_type_k": "--cache-type-k",
        "cache_type_v": "--cache-type-v",
        "flash_attn": "--flash-attn",
    }
    for key, value in args.items():
        flag = flag_map.get(key)
        if flag:
            if flag in cmd:
                continue
            _set_arg(cmd, flag, value)
    return cmd


def _merge_extra_args(cmd: list, extra: str) -> list:
    if not extra:
        return cmd
    extra_list = shlex.split(extra)
    filtered: list[str] = []
    skip_next = False
    for item in extra_list:
        if skip_next:
            skip_next = False
            continue
        if item in {"--device", "-dev"}:
            log.warning("Dropping --device from extra args to avoid llama.cpp device errors.")
            skip_next = True
            continue
        filtered.append(item)
    for flag in filtered:
        if flag not in cmd:
            cmd.append(flag)
    return cmd


def _update_model_command(command: Any, model_path: str, args: Dict[str, str], extra: str) -> list:
    cmd = _command_to_list(command)
    if "--device" in cmd:
        idx = cmd.index("--device")
        del cmd[idx: idx + 2]
    cmd = _set_arg(cmd, "--model", model_path)
    cmd = _merge_args(cmd, args)
    cmd = _merge_extra_args(cmd, extra)
    return cmd


def _replace_flags(cmd: list, flags: Dict[str, Optional[str]], extra: str) -> list:
    result = list(cmd)
    for flag in flags.keys():
        while flag in result:
            idx = result.index(flag)
            del result[idx: idx + 2]
    if "--device" in result:
        idx = result.index("--device")
        del result[idx: idx + 2]
    for flag, value in flags.items():
        if value is not None and value != "":
            result = _set_arg(result, flag, value)
    result = _merge_extra_args(result, extra)
    return result


async def get_app_config(cfg: TrueNASConfig) -> Dict[str, Any]:
    config = await _rpc_call(cfg, "app.config", [cfg.app_name])
    if not isinstance(config, dict):
        raise RuntimeError("app.config returned unsupported payload")
    return config


async def get_app_command(cfg: TrueNASConfig, service_name: str = "llamacpp") -> list:
    config = await get_app_config(cfg)
    return _extract_command(config, service_name=service_name)


async def get_active_model_id(cfg: TrueNASConfig, service_name: str = "llamacpp") -> str:
    config = await get_app_config(cfg)
    cmd = _extract_command(config, service_name=service_name)
    return _model_id_from_command(cmd) or ""


async def get_app_logs(
    cfg: TrueNASConfig,
    tail_lines: int = 200,
    service_name: str = "llamacpp",
) -> str:
    tail_payloads = [
        {"tail": tail_lines},
        {"tail_lines": tail_lines},
        {"tail": str(tail_lines)},
    ]
    for payload in tail_payloads:
        try:
            result = await _rpc_call(cfg, "app.container_logs", [cfg.app_name, service_name, payload])
            if isinstance(result, str):
                return result
        except Exception as exc:
            log.debug("app.container_logs failed (%s): %s", payload, exc)
    for payload in tail_payloads:
        try:
            result = await _rpc_call(cfg, "app.logs", [cfg.app_name, payload])
            if isinstance(result, str):
                return result
        except Exception as exc:
            log.debug("app.logs failed (%s): %s", payload, exc)
    return ""


async def update_app_command(
    cfg: TrueNASConfig,
    command: list,
    service_name: str = "llamacpp",
) -> None:
    config = await _rpc_call(cfg, "app.config", [cfg.app_name])
    if not isinstance(config, dict):
        raise RuntimeError("app.config returned unsupported payload")
    if config.get("custom_compose_config") or config.get("custom_compose_config_string"):
        compose = _parse_compose(config.get("custom_compose_config") or config.get("custom_compose_config_string") or {})
        services = compose.get("services") or {}
        if service_name not in services:
            raise RuntimeError(f"service {service_name} not found in compose")
        svc = services[service_name]
        svc["command"] = command
        await _rpc_call(cfg, "app.update", [cfg.app_name, {"custom_compose_config": compose}])
        return
    config["command"] = command
    await _rpc_call(cfg, "app.update", [cfg.app_name, {"values": config}])


async def update_command_flags(
    cfg: TrueNASConfig,
    flags: Dict[str, Optional[str]],
    extra: str,
    service_name: str = "llamacpp",
) -> None:
    config = await _rpc_call(cfg, "app.config", [cfg.app_name])
    if not isinstance(config, dict):
        raise RuntimeError("app.config returned unsupported payload")
    if config.get("custom_compose_config") or config.get("custom_compose_config_string"):
        compose = _parse_compose(config.get("custom_compose_config") or config.get("custom_compose_config_string") or {})
        services = compose.get("services") or {}
        if service_name not in services:
            raise RuntimeError(f"service {service_name} not found in compose")
        svc = services[service_name]
        cmd = svc.get("command")
        svc["command"] = _replace_flags(_command_to_list(cmd), flags, extra)
        await _rpc_call(cfg, "app.update", [cfg.app_name, {"custom_compose_config": compose}])
        return
    cmd = _replace_flags(_command_to_list(config.get("command")), flags, extra)
    config["command"] = cmd
    await _rpc_call(cfg, "app.update", [cfg.app_name, {"values": config}])


async def _rpc_call(cfg: TrueNASConfig, method: str, params: Optional[list] = None) -> Any:
    ssl_ctx = None
    if cfg.ws_url.startswith("wss://") and not cfg.verify_ssl:
        ssl_ctx = ssl.create_default_context()
        ssl_ctx.check_hostname = False
        ssl_ctx.verify_mode = ssl.CERT_NONE

    async with websockets.connect(cfg.ws_url, ssl=ssl_ctx) as ws:
        await ws.send(json.dumps({"msg": "connect", "version": "1", "support": ["1"]}))
        connected = json.loads(await ws.recv())
        if connected.get("msg") != "connected":
            raise RuntimeError("failed to connect to TrueNAS websocket")

        await ws.send(
            json.dumps({"id": 1, "msg": "method", "method": "auth.login_with_api_key", "params": [cfg.api_key]})
        )
        auth_resp = json.loads(await ws.recv())
        if not auth_resp.get("result"):
            if not cfg.api_user:
                raise RuntimeError("API key rejected and TRUENAS_API_USER not set")
            await ws.send(
                json.dumps(
                    {
                        "id": 2,
                        "msg": "method",
                        "method": "auth.login_ex",
                        "params": [
                            {
                                "mechanism": "API_KEY_PLAIN",
                                "username": cfg.api_user,
                                "api_key": cfg.api_key,
                            }
                        ],
                    }
                )
            )
            auth_ex = json.loads(await ws.recv())
            if auth_ex.get("result", {}).get("response_type") != "SUCCESS":
                raise RuntimeError("API key authentication failed")

        req_id = 3
        await ws.send(json.dumps({"id": req_id, "msg": "method", "method": method, "params": params or []}))
        while True:
            raw = json.loads(await ws.recv())
            if raw.get("id") != req_id:
                continue
            if raw.get("msg") == "error":
                raise RuntimeError(raw.get("error"))
            return raw.get("result")


async def switch_model(
    cfg: TrueNASConfig,
    model_path: str,
    args: Dict[str, str],
    extra: str,
    service_name: str = "llamacpp",
) -> None:
    config = await _rpc_call(cfg, "app.config", [cfg.app_name])
    if config.get("custom_compose_config") or config.get("custom_compose_config_string"):
        compose = _parse_compose(config.get("custom_compose_config") or config.get("custom_compose_config_string") or {})
        services = compose.get("services") or {}
        if service_name not in services:
            raise RuntimeError(f"service {service_name} not found in compose")
        svc = services[service_name]
        cmd = svc.get("command")
        svc["command"] = _update_model_command(cmd, model_path, args, extra)
        await _rpc_call(cfg, "app.update", [cfg.app_name, {"custom_compose_config": compose}])
        log.info("Requested model switch to %s via TrueNAS middleware (custom app)", model_path)
        return

    if not isinstance(config, dict):
        raise RuntimeError("app.config returned unsupported payload")

    cmd = config.get("command")
    config["command"] = _update_model_command(cmd, model_path, args, extra)
    await _rpc_call(cfg, "app.update", [cfg.app_name, {"values": config}])
    log.info("Requested model switch to %s via TrueNAS middleware (catalog app)", model_path)