"""
SSHService — asyncssh wrapper with per-device/per-object/global concurrency limits.

Each call to run_command() opens a fresh connection (no pooling).
Semaphores prevent overwhelming devices or the central server.

Auth resolution order for a device:
  1. device.ssh_last_auth  (last successful method — tried first)
  2. device.ssh_key_id_override  (specific SSH key for this device)
  3. device.ssh_pass_override_enc  (legacy password override for this device)
  4. object_category_credentials WHERE category=device.category AND protocol='ssh'
  5. object.ssh_credentials  (ordered credential list)
  6. object.ssh_user + object.ssh_pass_enc  (legacy fallback)

On success the working method is persisted to device.ssh_last_auth.
If the last_auth method fails it is cleared and the full list is tried.
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Optional

import asyncssh
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.device import Device
from app.models.object import Object
from app.models.ssh_key import SSHKey
from app.services.crypto import decrypt
from app.services.database import AsyncSessionLocal

logger = logging.getLogger(__name__)

# Concurrency limits
_GLOBAL_LIMIT = 50
_OBJECT_LIMIT = 10
_DEVICE_LIMIT = 1

_global_sem = asyncio.Semaphore(_GLOBAL_LIMIT)
_object_sems: dict[int, asyncio.Semaphore] = {}
_device_sems: dict[int, asyncio.Semaphore] = {}


def _get_object_sem(object_id: int) -> asyncio.Semaphore:
    if object_id not in _object_sems:
        _object_sems[object_id] = asyncio.Semaphore(_OBJECT_LIMIT)
    return _object_sems[object_id]


def _get_device_sem(device_id: int) -> asyncio.Semaphore:
    if device_id not in _device_sems:
        _device_sems[device_id] = asyncio.Semaphore(_DEVICE_LIMIT)
    return _device_sems[device_id]


@dataclass
class SSHResult:
    success: bool
    stdout: str = ""
    stderr: str = ""
    exit_code: int = -1
    duration_ms: int = 0
    error: str = ""
    # True when SSH auth succeeded (connection established), even if command returned non-zero.
    # Used to distinguish "wrong password" from "command failed".
    auth_ok: bool = False


@dataclass
class _AuthOption:
    """Internal representation of one auth attempt."""
    type: str           # "key" or "password"
    username: str
    password: str = ""  # for type=password
    key_path: str = ""  # for type=key
    # Identifiers used to persist last_auth
    cred_id: Optional[int] = None   # ObjectSSHCredential.id (if applicable)
    key_id: Optional[int] = None    # SSHKey.id (if applicable)


def _opt_key(opt: _AuthOption) -> str:
    """Deduplication key for an auth option."""
    if opt.type == "key":
        return f"key:{opt.key_id}:{opt.username}"
    return f"password:{opt.cred_id}:{opt.username}"


async def _load_key_path(db: AsyncSession, key_id: int) -> Optional[str]:
    """Return the filesystem path for an SSHKey id, or None if not found."""
    from sqlalchemy import select
    result = await db.execute(select(SSHKey).where(SSHKey.id == key_id))
    key = result.scalar_one_or_none()
    return key.path if key else None


async def _add_global_credentials(
    options: list[_AuthOption],
    seen: set[str],
    db: Optional[AsyncSession],
) -> None:
    """Append global SSH credentials (app-wide fallback) to the options list."""
    if not db:
        return
    from sqlalchemy import select as sa_select
    from app.models.ssh_key import GlobalSSHCredential

    def _add(opt: _AuthOption) -> None:
        k = _opt_key(opt)
        if k not in seen:
            seen.add(k)
            options.append(opt)

    result = await db.execute(
        sa_select(GlobalSSHCredential).order_by(GlobalSSHCredential.priority, GlobalSSHCredential.id)
    )
    for gcred in result.scalars().all():
        if gcred.type == "key" and gcred.ssh_key_id:
            path = await _load_key_path(db, gcred.ssh_key_id)
            if path:
                _add(_AuthOption(type="key", username=gcred.username, key_path=path, key_id=gcred.ssh_key_id))
        elif gcred.type == "password" and gcred.password_enc:
            _add(_AuthOption(type="password", username=gcred.username, password=decrypt(gcred.password_enc)))


async def build_auth_options(device: Device, obj: Object, db: Optional[AsyncSession] = None) -> list[_AuthOption]:
    """Build an ordered list of auth options to try for this device.

    If db is provided, key paths are resolved from the database.
    Without db, key-based options referencing the DB are skipped.
    """
    options: list[_AuthOption] = []
    seen: set[str] = set()

    def _add(opt: _AuthOption) -> None:
        k = _opt_key(opt)
        if k not in seen:
            seen.add(k)
            options.append(opt)

    port = device.ssh_port_override or obj.ssh_port or 22  # not used in options, used by caller

    # 1. Last successful method — try first
    if device.ssh_last_auth:
        la = device.ssh_last_auth
        if la.get("type") == "key" and la.get("key_id") and db:
            path = await _load_key_path(db, la["key_id"])
            if path:
                _add(_AuthOption(
                    type="key",
                    username=la.get("username", obj.ssh_user or "root"),
                    key_path=path,
                    key_id=la["key_id"],
                ))
        elif la.get("type") == "password":
            # We store the password reference; reconstruct from device/object
            username = la.get("username", obj.ssh_user or "root")
            cred_id = la.get("cred_id")
            # Try to get password from matching cred
            password = ""
            if cred_id:
                for cred in obj.ssh_credentials:
                    if cred.id == cred_id and cred.password_enc:
                        password = decrypt(cred.password_enc)
                        break
            if not password:
                if device.ssh_pass_override_enc:
                    password = decrypt(device.ssh_pass_override_enc)
                elif obj.ssh_pass_enc:
                    password = decrypt(obj.ssh_pass_enc)
            _add(_AuthOption(type="password", username=username, password=password, cred_id=cred_id))

    # 2. Device-specific SSH key override
    if device.ssh_key_id_override and db:
        path = await _load_key_path(db, device.ssh_key_id_override)
        if path:
            username = device.ssh_user_override or obj.ssh_user or "root"
            _add(_AuthOption(type="key", username=username, key_path=path, key_id=device.ssh_key_id_override))

    # 3. Device-specific password override (legacy)
    if device.ssh_pass_override_enc:
        username = device.ssh_user_override or obj.ssh_user or "root"
        password = decrypt(device.ssh_pass_override_enc)
        _add(_AuthOption(type="password", username=username, password=password))

    # 4. Per-category SSH credentials from object
    for cat_cred in obj.category_credentials:
        if cat_cred.category != device.category or cat_cred.protocol != "ssh":
            continue
        if cat_cred.type == "key" and cat_cred.ssh_key_id and db:
            path = await _load_key_path(db, cat_cred.ssh_key_id)
            if path:
                _add(_AuthOption(type="key", username=cat_cred.username, key_path=path, key_id=cat_cred.ssh_key_id))
        elif cat_cred.type == "password" and cat_cred.password_enc:
            _add(_AuthOption(
                type="password",
                username=cat_cred.username,
                password=decrypt(cat_cred.password_enc),
                cred_id=cat_cred.id,
            ))

    # 5. Object-level SSH credential list
    for cred in obj.ssh_credentials:
        if cred.type == "key" and cred.ssh_key_id and db:
            path = await _load_key_path(db, cred.ssh_key_id)
            if path:
                _add(_AuthOption(type="key", username=cred.username, key_path=path, key_id=cred.ssh_key_id))
        elif cred.type == "password" and cred.password_enc:
            _add(_AuthOption(
                type="password",
                username=cred.username,
                password=decrypt(cred.password_enc),
                cred_id=cred.id,
            ))

    # 6. Legacy object ssh_user + ssh_pass_enc fallback
    if obj.ssh_user and obj.ssh_pass_enc:
        _add(_AuthOption(
            type="password",
            username=obj.ssh_user,
            password=decrypt(obj.ssh_pass_enc),
        ))

    # 7. Global app-wide credentials (last resort)
    await _add_global_credentials(options, seen, db)

    return options


async def _try_connect(
    host: str,
    port: int,
    opt: _AuthOption,
    command: str,
    timeout: int,
) -> SSHResult:
    """Attempt a single SSH connection with the given auth option."""
    start = time.monotonic()
    connect_kwargs: dict = dict(
        host=host,
        port=port,
        username=opt.username,
        known_hosts=None,
        connect_timeout=10,
    )
    if opt.type == "key":
        connect_kwargs["client_keys"] = [opt.key_path]
        connect_kwargs["passphrase"] = None
    else:
        connect_kwargs["password"] = opt.password

    try:
        conn = await asyncio.wait_for(asyncssh.connect(**connect_kwargs), timeout=15)
    except asyncio.TimeoutError:
        return SSHResult(
            success=False,
            error="SSH connection timed out",
            duration_ms=int((time.monotonic() - start) * 1000),
        )
    except (asyncssh.DisconnectError, asyncssh.PermissionDenied, asyncssh.ConnectionLost) as e:
        return SSHResult(
            success=False,
            error=str(e),
            duration_ms=int((time.monotonic() - start) * 1000),
        )
    except Exception as e:
        return SSHResult(
            success=False,
            error=str(e),
            duration_ms=int((time.monotonic() - start) * 1000),
        )

    try:
        async with conn:
            result = await asyncio.wait_for(conn.run(command, check=False), timeout=timeout)
        duration = int((time.monotonic() - start) * 1000)
        return SSHResult(
            success=result.exit_status == 0,
            stdout=(result.stdout or "").strip(),
            stderr=(result.stderr or "").strip(),
            exit_code=result.exit_status or 0,
            duration_ms=duration,
            auth_ok=True,
        )
    except asyncio.TimeoutError:
        return SSHResult(
            success=False,
            error=f"SSH command timed out after {timeout}s",
            duration_ms=int((time.monotonic() - start) * 1000),
            auth_ok=True,
        )
    except Exception as e:
        return SSHResult(
            success=False,
            error=str(e),
            duration_ms=int((time.monotonic() - start) * 1000),
            auth_ok=True,
        )

async def _save_last_auth(device: Device, db: AsyncSession, opt: _AuthOption) -> None:
    """Persist the successful auth method to device.ssh_last_auth."""
    payload: dict = {"type": opt.type, "username": opt.username}
    if opt.type == "key":
        payload["key_id"] = opt.key_id
    else:
        if opt.cred_id:
            payload["cred_id"] = opt.cred_id
    device.ssh_last_auth = payload
    try:
        await db.flush()
    except Exception:
        pass  # non-critical


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

# Legacy helper kept for callers that don't have device/object context
def resolve_ssh_credentials(device: Device, obj: Object) -> tuple[str, str, int]:
    """Returns (username, password, port) — device overrides take priority.

    DEPRECATED: use run_command() with db= for full multi-auth support.
    """
    user = device.ssh_user_override or obj.ssh_user or "root"
    port = device.ssh_port_override or obj.ssh_port or 22

    if device.ssh_pass_override_enc:
        password = decrypt(device.ssh_pass_override_enc)
    elif obj.ssh_pass_enc:
        password = decrypt(obj.ssh_pass_enc)
    else:
        password = ""

    return user, password, port


async def run_command(
    device: Device,
    obj: Object,
    command: str,
    timeout: int = 30,
    db: Optional[AsyncSession] = None,
) -> SSHResult:
    """Run command on device, trying all configured auth methods in order.

    If db is provided:
    - SSH keys are resolved from the database.
    - The successful auth method is persisted to device.ssh_last_auth.

    If last_auth method fails, it is cleared and remaining methods are tried.
    """
    port = device.ssh_port_override or obj.ssh_port or 22

    async with _global_sem:
        async with _get_object_sem(device.object_id):
            async with _get_device_sem(device.id):
                if db:
                    options = await build_auth_options(device, obj, db)
                else:
                    async with AsyncSessionLocal() as auth_db:
                        options = await build_auth_options(device, obj, auth_db)

                if not options:
                    # Fallback: no creds configured at all
                    user = device.ssh_user_override or obj.ssh_user or "root"
                    options = [_AuthOption(type="password", username=user, password="")]

                last_auth_option = options[0] if device.ssh_last_auth else None
                last_result: Optional[SSHResult] = None

                for i, opt in enumerate(options):
                    if opt.type == "key":
                        logger.info(
                            "SSH attempt %d/%d for %s — key: user=%s key_path=%s",
                            i + 1, len(options), device.ip, opt.username, opt.key_path,
                        )
                    else:
                        masked = (opt.password[:1] + "***" + opt.password[-1:]) if len(opt.password) > 2 else "***"
                        logger.info(
                            "SSH attempt %d/%d for %s — password: user=%s password=%s",
                            i + 1, len(options), device.ip, opt.username, masked,
                        )

                    result = await _try_connect(device.ip, port, opt, command, timeout)

                    if result.auth_ok:
                        # Connection was established — auth succeeded.
                        # Return the result regardless of exit_code (command may have failed, not auth).
                        logger.info(
                            "SSH auth succeeded for %s with %s(%s)",
                            device.ip, opt.type, opt.username,
                        )
                        if db:
                            await _save_last_auth(device, db, opt)
                        else:
                            async with AsyncSessionLocal() as save_db:
                                async with save_db.begin():
                                    current_device = await save_db.get(Device, device.id)
                                    if current_device:
                                        await _save_last_auth(current_device, save_db, opt)
                        return result

                    # Auth failed (wrong password / key / connection error)
                    # If last_auth method failed, clear it
                    if i == 0 and opt is last_auth_option:
                        device.ssh_last_auth = None
                        if db:
                            try:
                                await db.flush()
                            except Exception:
                                pass
                        else:
                            async with AsyncSessionLocal() as clear_db:
                                async with clear_db.begin():
                                    current_device = await clear_db.get(Device, device.id)
                                    if current_device:
                                        current_device.ssh_last_auth = None

                    last_result = result
                    logger.info(
                        "SSH auth failed for %s with %s(%s): %s",
                        device.ip, opt.type, opt.username, result.error,
                    )

                # All methods exhausted
                if last_result:
                    last_result.error = f"All auth methods exhausted. Last error: {last_result.error}"
                    return last_result
                return SSHResult(success=False, error="No auth methods available")


async def run_command_direct(
    host: str,
    port: int,
    user: str,
    password: str,
    command: str,
    timeout: int = 30,
) -> SSHResult:
    """Connect directly with explicit credentials (for object-level hosts like rack PCs).

    Used by device_discovery when connecting to a specific IP that isn't a Device record.
    """
    opt = _AuthOption(type="password", username=user, password=password)
    async with _global_sem:
        return await _try_connect(host, port, opt, command, timeout)


async def run_command_direct_with_options(
    host: str,
    port: int,
    options: list[_AuthOption],
    command: str,
    timeout: int = 30,
) -> SSHResult:
    """Try multiple auth options for a non-Device host (e.g. rack PC during discovery)."""
    async with _global_sem:
        last_result: Optional[SSHResult] = None
        for i, opt in enumerate(options):
            if opt.type == "key":
                logger.info(
                    "SSH attempt %d/%d for %s — key: user=%s key_path=%s",
                    i + 1, len(options), host, opt.username, opt.key_path,
                )
            else:
                masked = (opt.password[:1] + "***" + opt.password[-1:]) if len(opt.password) > 2 else "***"
                logger.info(
                    "SSH attempt %d/%d for %s — password: user=%s password=%s",
                    i + 1, len(options), host, opt.username, masked,
                )
            result = await _try_connect(host, port, opt, command, timeout)
            if result.auth_ok:
                logger.info("SSH auth succeeded for %s with %s(%s)", host, opt.type, opt.username)
                return result
            logger.info("SSH auth failed for %s with %s(%s): %s", host, opt.type, opt.username, result.error)
            last_result = result
        if last_result:
            last_result.error = f"All auth methods exhausted. Last error: {last_result.error}"
            return last_result
        return SSHResult(success=False, error="No auth methods available")


async def build_object_auth_options(obj: Object, db: Optional[AsyncSession] = None) -> list[_AuthOption]:
    """Build auth options for connecting to object-level hosts (server_ip, rack PCs).

    Uses the object's credential list + legacy ssh_user/ssh_pass_enc.
    No device-level overrides.
    """
    options: list[_AuthOption] = []
    seen: set[str] = set()

    def _add(opt: _AuthOption) -> None:
        k = _opt_key(opt)
        if k not in seen:
            seen.add(k)
            options.append(opt)

    # Object SSH credential list
    for cred in obj.ssh_credentials:
        if cred.type == "key" and cred.ssh_key_id and db:
            path = await _load_key_path(db, cred.ssh_key_id)
            if path:
                _add(_AuthOption(type="key", username=cred.username, key_path=path, key_id=cred.ssh_key_id))
        elif cred.type == "password" and cred.password_enc:
            _add(_AuthOption(
                type="password",
                username=cred.username,
                password=decrypt(cred.password_enc),
                cred_id=cred.id,
            ))

    # Legacy fallback
    if obj.ssh_user and obj.ssh_pass_enc:
        _add(_AuthOption(type="password", username=obj.ssh_user, password=decrypt(obj.ssh_pass_enc)))

    # Global app-wide credentials (last resort)
    await _add_global_credentials(options, seen, db)

    return options