import asyncio
import asyncssh
import traceback
from fastapi import APIRouter, WebSocket, WebSocketDisconnect

router = APIRouter()

@router.websocket("/ws/terminal")
async def websocket_terminal(websocket: WebSocket, host: str, username: str, password: str):
    await websocket.accept()
    conn = 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: 
                        print("[DEBUG] 設備端主動關閉了 stdout 通道")
                        break
                    
                    safe_text = data_bytes.decode('utf-8', errors='replace')
                    await websocket.send_text(safe_text)
            except Exception as e:
                print("\n❌ [WS Forward Error] 讀取設備畫面時發生錯誤:")
                traceback.print_exc()

        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:
                pass
            except Exception as e:
                print("\n❌ [SSH Forward Error] 寫入指令到設備時發生錯誤:")
                traceback.print_exc()

        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 Exception as e:
        error_msg = f"\r\n\x1b[31mSSH Connection Error: {str(e)}\x1b[0m\r\n"
        try:
            await websocket.send_text(error_msg)
        except:
            pass
        print("\n❌ [Connection Error] 建立連線或執行過程中發生錯誤:")
        traceback.print_exc()
    finally:
        if conn:
            conn.close()
        try:
            await websocket.close()
        except:
            pass
