import asyncio
import asyncssh
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from logger import get_logger

logger = get_logger("app.ssh")
router = APIRouter()

@router.websocket("/ws/terminal")
async def websocket_terminal(websocket: WebSocket, host: str, username: str, password: str):
    await websocket.accept()
    conn = None
    process = None
    ws_task = None
    ssh_task = None
    try:
        # 建立連線
        conn = await asyncssh.connect(host, username=username, password=password, known_hosts=None)
        
        # 💡 關鍵修復：使用 term_size 參數來指定寬高 (width, height)
        process = await conn.create_process(
            term_type='xterm-256color', 
            term_size=(80, 24),  # 正確的 asyncssh 參數格式
            encoding=None
        )

        async def forward_to_ws():
            try:
                while True:
                    data_bytes = await process.stdout.read(8192)
                    if not data_bytes: 
                        logger.debug("設備端主動關閉了 stdout 通道")
                        break
                    
                    safe_text = data_bytes.decode('utf-8', errors='replace')
                    await websocket.send_text(safe_text)
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error("❌ [WS Forward Error] 讀取設備畫面時發生錯誤", exc_info=True)

        async def forward_to_ssh():
            try:
                while True:
                    data_text = await websocket.receive_text()
                    data_bytes = data_text.encode('utf-8')
                    
                    # 攔截 xterm.js 預設的 DEL (\x7f)，轉換為 BS (\x08)
                    if b'\x7f' in data_bytes:
                        data_bytes = data_bytes.replace(b'\x7f', b'\x08')
                        
                    process.stdin.write(data_bytes)
                    await process.stdin.drain()
            except WebSocketDisconnect:
                # 主動拋出，讓外層捕捉以進行資源回收
                # raise
                # 🌟 修正：不要再 raise 拋出去了，直接 return 結束任務，讓外層自然回收
                logger.debug("前端 WebSocket 正常斷開 (使用者重整或關閉網頁)")
                return
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error("❌ [SSH Forward Error] 寫入指令到設備時發生錯誤", exc_info=True)

        ws_task = asyncio.create_task(forward_to_ws())
        ssh_task = asyncio.create_task(forward_to_ssh())
        
        done, pending = await asyncio.wait(
            [ws_task, ssh_task],
            return_when=asyncio.FIRST_COMPLETED
        )
        
        for task in pending:
            task.cancel()

    except WebSocketDisconnect:
        logger.debug("WebSocket 正常斷線，正在終止背景任務...")
        if ws_task and not ws_task.done():
            ws_task.cancel()
        if ssh_task and not ssh_task.done():
            ssh_task.cancel()
    except Exception as e:
        err_str = str(e)
        error_msg = f"\r\n\x1b[31mSSH Connection Error: {err_str}\x1b[0m\r\n"
        try:
            # 先把錯誤印在終端機畫面上
            await websocket.send_text(error_msg)
            
            # 🌟 關鍵修復：使用自訂斷線碼 4001，並附上錯誤原因 (WebSocket 規範 reason 最長 123 bytes)
            safe_reason = err_str[:120] if err_str else "Authentication or Connection Failed"
            await websocket.close(code=4001, reason=safe_reason)
        except:
            pass
        logger.error("❌ [Connection Error] 建立連線或執行過程中發生錯誤", exc_info=True)
        return  # 🌟 提早結束，避免下方的 finally 再次執行 close 導致報錯
    finally:
        # 🌟 確保 process 與 conn 絕對被關閉，防止 Memory Leak
        if process:
            process.close()
        if conn:
            conn.close()
        try:
            await websocket.close()
        except:
            pass