import os
import asyncio
import json
import pty
import struct
import fcntl
import termios
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()
PASSWORD = os.getenv("TERMINAL_PASSWORD")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, password: str = Query(default="")):
    if password != PASSWORD:
        await websocket.close(code=1008)
        print("Conexão recusada: senha incorreta")
        return
    
    await websocket.accept()
    
    pid, fd = pty.fork()
    
    if pid == 0:
        # Processo filho
        os.environ["TERM"] = "xterm-256color"
        shell = os.environ.get("SHELL", "/bin/bash")
        os.execvp(shell, [shell])
    else:
        # Processo pai
        loop = asyncio.get_event_loop()
        
        async def read_from_terminal():
            while True:
                try:
                    data = await loop.run_in_executor(None, os.read, fd, 1024)
                    if data:
                        await websocket.send_text(data.decode(errors="ignore"))
                    else:
                        break
                except Exception as e:
                    print("Erro ao ler do terminal:", e)
                    break
        
        async def write_to_terminal():
            try:
                while True:
                    msg = await websocket.receive_text()
                    
                    # Tenta parsear como JSON
                    try:
                        data = json.loads(msg)
                        
                        # Handle resize
                        if data.get("type") == "resize":
                            cols = int(data.get("cols", 80))
                            rows = int(data.get("rows", 24))
                            
                            # Método mais seguro usando ioctl
                            try:
                                winsize = struct.pack("HHHH", rows, cols, 0, 0)
                                fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
                                print(f"Terminal redimensionado: {cols}x{rows}")
                            except Exception as e:
                                print(f"Erro ao redimensionar terminal: {e}")
                            
                            # Não envia nada para o terminal, apenas continua
                            continue
                            
                    except json.JSONDecodeError:
                        # Se não for JSON, é input normal do usuário
                        os.write(fd, msg.encode())
                    except ValueError:
                        # Se houver erro nos valores do resize
                        print(f"Valores inválidos no resize: {msg}")
                        continue
                        
            except WebSocketDisconnect:
                print("Cliente desconectado")
            except Exception as e:
                print(f"Erro no write_to_terminal: {e}")
            finally:
                try:
                    os.close(fd)
                    os.kill(pid, 9)
                except:
                    pass
        
        try:
            await asyncio.gather(read_from_terminal(), write_to_terminal())
        except Exception as e:
            print(f"Erro na conexão: {e}")
        finally:
            try:
                os.close(fd)
                os.kill(pid, 9)
            except:
                pass