"""Nexus — Web SSH Terminal WebSocket Endpoint
"""

import asyncio
import json
import logging
import uuid
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query, Path

from server.config import settings
from server.domain.models import SshSession, CommandLog, Server, AuditLog
from server.infrastructure.ssh.asyncssh_pool import ssh_pool
from server.infrastructure.database.session import AsyncSessionLocal
from server.infrastructure.database.admin_repo import AdminRepositoryImpl
from server.infrastructure.database.server_repo import ServerRepositoryImpl
from server.infrastructure.database.ssh_session_repo import SshSessionRepositoryImpl, CommandLogRepositoryImpl
from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl

logger = logging.getLogger("nexus.webssh")

router = APIRouter()

# Koko protocol message types
MSG_TERMINAL_INIT = "TERMINAL_INIT"
MSG_DATA = "DATA"
MSG_RESIZE = "RESIZE"
MSG_CLOSE = "CLOSE"
MSG_ERROR = "ERROR"
MSG_PING = "PING"
MSG_PONG = "PONG"

# Command logging configuration
COMMAND_LOG_MAX_LENGTH = 2000


async def _verify_webssh_token(token: str):
    """Verify JWT token for WebSSH — returns (admin, server_id) or (None, None)

    Checks: valid signature, not expired, sub exists, admin active, and
    updated_at matches (invalidates token after password change).
    """
    try:
        import jwt as pyjwt
        payload = pyjwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"],
                                options={"require": ["exp", "sub"]})
        if payload.get("purpose") != "webssh":
            return None, None

        admin_id = payload.get("sub")
        server_id = payload.get("server_id")
        if not admin_id or server_id is None:
            return None, None
        try:
            server_id = int(server_id)
        except (TypeError, ValueError):
            return None, None

        async with AsyncSessionLocal() as session:
            admin_repo = AdminRepositoryImpl(session)
            admin = await admin_repo.get_by_id(admin_id)
            if not admin or not admin.is_active:
                return None, None

            # P1-7: Check updated_at — invalidate token after password change
            token_updated = payload.get("updated", 0)
            if token_updated and admin.updated_at:
                admin_ts = int(admin.updated_at.timestamp())
                # Allow 5-second grace for clock skew / race conditions
                if admin_ts > token_updated + 5:
                    return None, None

            return admin, server_id
    except Exception:
        logger.debug("WebSSH JWT verification failed", exc_info=True)
        return None, None


@router.websocket("/ws/terminal/{server_id}")
async def terminal_ws(
    websocket: WebSocket,
    server_id: int = Path(..., description="Target server ID"),
    token: Optional[str] = Query(None, description="JWT access token"),
):
    """WebSSH Terminal WebSocket endpoint

    ADR-011: JWT authentication via query parameter `?token=xxx`
    W3: WebSocket endpoint for terminal access
    W5: Koko message protocol
    """
    # ── JWT Authentication ──
    if not token:
        await websocket.close(code=4001, reason="Missing JWT token")
        return

    admin, token_server_id = await _verify_webssh_token(token)
    if not admin:
        await websocket.close(code=4001, reason="Invalid or expired JWT token")
        return

    # Security: WebSSH token must bind to this server (general access_token rejected)
    if token_server_id != server_id:
        await websocket.close(code=4003, reason="Token not authorized for this server")
        return

    # ── P2-16: Consolidated pre-accept DB operations in one session ──
    # server lookup + audit log + SSH session record all share one session,
    # which is closed before the long-lived WebSocket relay begins.
    client_ip = websocket.client.host if websocket.client else "unknown"
    session_id = str(uuid.uuid4())
    server: Optional[Server] = None

    async with AsyncSessionLocal() as db:
        server_repo = ServerRepositoryImpl(db)
        server = await server_repo.get_by_id(server_id)

        if not server:
            await websocket.close(code=4003, reason="Server not found")
            return

        # ── Server-level authorization check ──
        if not server.domain:
            await websocket.close(code=4003, reason="Server domain not configured")
            return
        if server.auth_method == "key" and not server.ssh_key_path and not server.ssh_key_private:
            await websocket.close(code=4003, reason="SSH key not configured")
            return
        if server.auth_method == "password" and not server.password:
            await websocket.close(code=4003, reason="SSH password not configured")
            return

        # Audit: WebSSH session start
        audit_repo = AuditLogRepositoryImpl(db)
        await audit_repo.create(AuditLog(
            admin_username=admin.username,
            action="webssh_connect",
            target_type="server",
            target_id=server.id,
            detail=f"WebSSH 连接 {server.name} ({server.domain})",
            ip_address=client_ip,
        ))

        # Create SSH session record
        ssh_session_repo = SshSessionRepositoryImpl(db)
        ssh_session = SshSession(
            id=session_id,
            server_id=server.id,
            admin_id=admin.id,
            remote_addr=client_ip,
            status="active",
        )
        await ssh_session_repo.create(ssh_session)

    # ── Accept WebSocket ──
    await websocket.accept()

    logger.info(f"WebSSH: admin={admin.username} connecting to server={server.name} ({server.domain})")

    # ── Establish SSH connection (with stale-connection retry) ──
    conn = None
    try:
        conn = await ssh_pool.acquire(server)
    except Exception as e:
        logger.error(f"WebSSH connection failed: {e}")
        try:
            await websocket.send_json({"type": MSG_ERROR, "message": f"SSH连接失败: {str(e)}"})
        except Exception:  # noqa: S110 — best-effort cleanup
            pass
        try:
            await websocket.close(code=4003, reason="SSH connection failed")
        except Exception:  # noqa: S110 — best-effort cleanup
            pass
        await _close_ssh_session(session_id)
        return

    # ── Create SSH shell (retry once if stale pooled connection) ──
    ws_closed = False  # Track whether we've closed the WebSocket
    try:
        try:
            shell_proc = conn.create_process(
                term_type="xterm-256color",
                term_size=(24, 80),
                bufsize=4096,  # Small buffer for low-latency interactive echo
            )
        except Exception as e:
            # Stale pooled connection — force-close and retry with fresh connection
            logger.warning(f"WebSSH shell creation failed on pooled connection (server={server.id}): {type(e).__name__}: {e!r}")
            await ssh_pool.close_connection(server.id)
            try:
                conn = await ssh_pool.acquire(server)
                shell_proc = conn.create_process(
                    term_type="xterm-256color",
                    term_size=(24, 80),
                    bufsize=4096,  # Small buffer for low-latency interactive echo
                )
            except Exception as e2:
                logger.error(f"WebSSH shell creation failed on fresh connection (server={server.id}): {type(e2).__name__}: {e2!r}")
                try:
                    await websocket.send_json({"type": MSG_ERROR, "message": f"Shell创建失败: {type(e2).__name__}"})
                except Exception:  # noqa: S110 — best-effort cleanup
                    pass
                try:
                    await websocket.close(code=4003, reason="Shell creation failed")
                except Exception:  # noqa: S110 — best-effort cleanup
                    pass
                ws_closed = True
                await _close_ssh_session(session_id)
                return

        async with shell_proc as shell:
            # Send TERMINAL_INIT to client
            await websocket.send_json({
                "type": MSG_TERMINAL_INIT,
                "session_id": session_id,
                "server_name": server.name,
                "server_id": server.id,
            })

            # Color + CWD: ANSI/256/truecolor for ls, grep, vim, etc.
            shell.stdin.write(
                "export TERM=xterm-256color COLORTERM=truecolor "
                "FORCE_COLOR=1 CLICOLOR_FORCE=1\n"
            )
            shell.stdin.write('export PROMPT_COMMAND=\'echo -ne "\\033]0;${PWD}\\a"\'\n')

            # ── Bidirectional relay ──
            command_buffer = ""  # Buffer for command logging

            async def _read_shell_output():
                """Read SSH shell output -> send to WebSocket client

                Uses read() instead of async-for to avoid readline() buffering.
                async-for calls readline() internally, which waits for newline
                characters before yielding — causing echo delay in interactive
                terminal use.  read(n) returns data as soon as any is available.
                """
                try:
                    while not shell.stdout.at_eof():
                        data = await shell.stdout.read(4096)
                        if data:
                            await websocket.send_json({
                                "type": MSG_DATA,
                                "data": data,
                            })
                except Exception:  # noqa: S110 — SSH read loop, connection teardown
                    pass

            async def _read_websocket_input():
                """Read WebSocket client input -> write to SSH shell"""
                nonlocal command_buffer
                try:
                    while True:
                        raw = await websocket.receive_text()
                        try:
                            msg = json.loads(raw)
                        except json.JSONDecodeError:
                            # Raw text input (backward compat)
                            shell.stdin.write(raw)
                            command_buffer += raw
                            continue

                        msg_type = msg.get("type", "")

                        if msg_type == MSG_DATA:
                            data = msg.get("data", "")
                            shell.stdin.write(data)

                            # Buffer for command logging
                            command_buffer += data

                            # Log complete commands (on Enter)
                            if "\r" in data or "\n" in data:
                                cmd = command_buffer.strip()
                                if cmd:
                                    # P2-18: Check for dangerous commands
                                    from server.api.dependencies import check_dangerous_command
                                    check_dangerous_command(cmd)

                                    await _log_command(
                                        session_id=session_id,
                                        server_id=server.id,
                                        admin_id=admin.id,
                                        command=cmd[:COMMAND_LOG_MAX_LENGTH],
                                        remote_addr=client_ip,
                                    )
                                command_buffer = ""

                        elif msg_type == MSG_RESIZE:
                            cols = msg.get("cols", 80)
                            rows = msg.get("rows", 24)
                            shell.change_terminal_size(rows, cols)

                        elif msg_type == MSG_PING:
                            await websocket.send_json({"type": MSG_PONG, "ts": msg.get("ts")})

                        elif msg_type == MSG_CLOSE:
                            break

                except WebSocketDisconnect:
                    pass
                except Exception:  # noqa: S110 — WebSocket read loop, connection teardown
                    pass

            # Run both directions concurrently
            read_task = asyncio.create_task(_read_shell_output())
            write_task = asyncio.create_task(_read_websocket_input())

            try:
                # Wait for either direction to finish
                done, pending = await asyncio.wait(
                    [read_task, write_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for task in pending:
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass
            except Exception as e:
                logger.warning(f"WebSSH relay error: {e}")

    except Exception as e:
        logger.error(f"WebSSH shell creation failed: {type(e).__name__}: {e!r}")
        if not ws_closed:
            try:
                await websocket.send_json({"type": MSG_ERROR, "message": f"Shell创建失败: {type(e).__name__}"})
            except Exception:  # noqa: S110 — best-effort cleanup
                pass
    finally:
        # ── Cleanup ──
        await ssh_pool.release(server.id)
        await _close_ssh_session(session_id)

        if not ws_closed:
            try:
                await websocket.send_json({"type": MSG_CLOSE, "session_id": session_id})
            except Exception:  # noqa: S110 — best-effort cleanup
                pass

            try:
                await websocket.close()
            except Exception:  # noqa: S110 — best-effort cleanup
                pass

        logger.info(f"WebSSH session closed: admin={admin.username}, server={server.name}, session={session_id}")

        # Audit: WebSSH session end
        try:
            async with AsyncSessionLocal() as audit_session:
                audit_repo = AuditLogRepositoryImpl(audit_session)
                await audit_repo.create(AuditLog(
                    admin_username=admin.username,
                    action="webssh_disconnect",
                    target_type="server",
                    target_id=server.id,
                    detail=f"WebSSH 断开 {server.name}",
                    ip_address=client_ip,
                ))
        except Exception:
            logger.warning(f"Failed to write disconnect audit log for session {session_id}")


async def _close_ssh_session(session_id: str):
    """Close SSH session record in database"""
    try:
        async with AsyncSessionLocal() as session:
            ssh_session_repo = SshSessionRepositoryImpl(session)
            await ssh_session_repo.close_session(session_id)
    except Exception as e:
        logger.error(f"Failed to close SSH session {session_id}: {e}")


async def _log_command(session_id: str, server_id: int, admin_id: int, command: str, remote_addr: str):
    """Log a command to CommandLog

    Uses its own session per command — the WebSocket connection is long-lived
    (potentially hours), so we can't hold a single DB session open for the
    entire duration. Each command gets its own short-lived session.
    """
    try:
        async with AsyncSessionLocal() as session:
            cmd_log_repo = CommandLogRepositoryImpl(session)
            log = CommandLog(
                session_id=session_id,
                server_id=server_id,
                admin_id=admin_id,
                command=command,
                remote_addr=remote_addr,
            )
            await cmd_log_repo.create(log)
    except Exception as e:
        logger.error(f"Failed to log command: {e}")
