"""
HTTP + WebSocket Remote Shell Server
SSH Replacement với FastAPI
"""
import asyncio
import os
import signal
import pty
import subprocess
import io
import zipfile
import struct
import fcntl
import termios
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query, HTTPException, Depends, UploadFile, File as FastAPIFile
from fastapi.responses import StreamingResponse, FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI(title="Remote Shell")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# =============================================================================
# Configuration
# =============================================================================
# Use a more secure default if no token provided
SECRET_TOKEN = os.getenv("API_TOKEN")
if not SECRET_TOKEN:
    import secrets
    SECRET_TOKEN = secrets.token_urlsafe(32)
    print(f"⚠️  No API_TOKEN found in env. Generated temporary token: {SECRET_TOKEN}")

WORK_DIR = Path(os.getenv("WORK_DIR", "/tmp/remote-shell")).resolve()
WORK_DIR.mkdir(parents=True, exist_ok=True)

security = HTTPBearer(auto_error=False)

# Store running processes
running_processes: dict = {}
pty_sessions: dict = {}


# =============================================================================
# Auth Middleware
# =============================================================================
async def verify_token(credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)):
    """Verify Bearer token"""
    if credentials is None or credentials.credentials != SECRET_TOKEN:
        raise HTTPException(status_code=401, detail="Invalid or missing token")
    return credentials.credentials


def get_safe_path(path: str, session_id: str = None) -> Path:
    """Resolve path, supporting ~ for home directory. If session_id is provided, resolve relative to session workspace."""
    # Base directory depends on session_id
    base = WORK_DIR
    if session_id:
        # Sanitize session_id to prevent ".." in session_id
        safe_session_id = "".join(c for c in session_id if c.isalnum() or c in ("-", "_"))
        base = (WORK_DIR / safe_session_id).resolve()
        base.mkdir(parents=True, exist_ok=True)
    
    # Expand ~ to home directory (or base dir if in session)
    if path.startswith('~'):
        if session_id:
            path = path.replace('~', '.', 1)
        else:
            path = os.path.expanduser(path)
    
    # Resolve the path relative to base
    # This correctly handles both absolute paths (bypassing base if not careful) 
    # and relative paths.
    requested_path = Path(path)
    
    if requested_path.is_absolute():
        # If the server is in 'Global' mode (no session), allow absolute paths
        # If in 'Session' mode, absolute paths MUST be within the session base
        full_path = requested_path.resolve()
    else:
        full_path = (base / path).resolve()
    
    # CRITICAL: Path Traversal Check
    # Ensure the resolved path is still inside the intended base
    if not str(full_path).startswith(str(base)):
        raise HTTPException(status_code=403, detail=f"Access denied: path {path} is outside the allowed workspace")
    
    return full_path


# =============================================================================
# Command Execution (HTTP Streaming)
# =============================================================================
async def run_command_generator(cmd: str, session_id: str):
    """Generator to stream command output"""
    process = await asyncio.create_subprocess_shell(
        cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        cwd=str(WORK_DIR),
        preexec_fn=os.setsid
    )
    running_processes[session_id] = process

    try:
        while True:
            line = await process.stdout.readline()
            if not line:
                break
            yield line.decode('utf-8', errors='replace')
    finally:
        await process.wait()
        running_processes.pop(session_id, None)


@app.get("/exec")
async def stream_exec(
    cmd: str,
    session_id: str = Query(default="default"),
    token: str = Depends(verify_token)
):
    """Execute command with streaming output"""
    return StreamingResponse(
        run_command_generator(cmd, session_id),
        media_type="text/plain"
    )


# =============================================================================
# File Operations
# =============================================================================
@app.get("/files")
async def list_files(
    path: str = Query(default="."),
    show_hidden: bool = Query(default=True),
    session_id: str = Query(default=None),
    token: str = Depends(verify_token)
):
    """List directory contents including hidden files"""
    target = get_safe_path(path, session_id)
    if not target.exists():
        raise HTTPException(404, "Path not found")
    
    if target.is_file():
        stat = target.stat()
        return [{
            "name": target.name,
            "type": "file",
            "size": stat.st_size
        }]
    
    items = []
    for item in target.iterdir():
        # Skip hidden files if show_hidden is False
        if not show_hidden and item.name.startswith('.'):
            continue
        try:
            stat = item.stat()
            items.append({
                "name": item.name,
                "type": "directory" if item.is_dir() else "file",
                "size": stat.st_size if item.is_file() else None,
                "hidden": item.name.startswith('.')
            })
        except (PermissionError, OSError):
            # Skip files we can't access
            continue
    
    # Sort: directories first, then by name
    return sorted(items, key=lambda x: (x["type"] == "file", x["name"].lower()))


@app.get("/files/read")
async def read_file(
    path: str,
    session_id: str = Query(default=None),
    token: str = Depends(verify_token)
):
    """Read file content"""
    target = get_safe_path(path, session_id)
    if not target.exists():
        raise HTTPException(404, "File not found")
    if target.is_dir():
        raise HTTPException(400, "Cannot read directory")
    
    try:
        content = target.read_text(errors='replace')
        return {"path": path, "content": content}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/files/write")
async def write_file(
    path: str,
    content: str,
    session_id: str = Query(default=None),
    token: str = Depends(verify_token)
):
    """Write content to file"""
    target = get_safe_path(path, session_id)
    target.parent.mkdir(parents=True, exist_ok=True)
    
    try:
        target.write_text(content)
        return {"status": "ok", "path": path, "size": len(content)}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/files/upload-zip")
async def upload_zip_form(
    file: UploadFile = FastAPIFile(...),
    path: str = Query(default="."),
    session_id: str = Query(default=None),
    token: str = Depends(verify_token)
):
    """Upload and extract a zip file (multipart form).
    
    Args:
        file: Uploaded zip file
        path: Directory to extract to
        session_id: Session identifier for workspace isolation
    
    Returns:
        folder_path, tree
    """
    try:
        # Read uploaded file
        zip_bytes = await file.read()
        zip_filename = file.filename or "upload.zip"
        
        # Determine extract location
        target_dir = get_safe_path(path, session_id)
        target_dir.mkdir(parents=True, exist_ok=True)
        
        # Remove .zip extension for folder name
        folder_name = zip_filename.rsplit('.', 1)[0] if '.' in zip_filename else zip_filename
        folder_name = folder_name.replace(' ', '_').replace('/', '_')
        
        extract_dir = target_dir / folder_name
        extract_dir.mkdir(parents=True, exist_ok=True)
        
        # Extract zip
        try:
            with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
                zf.extractall(extract_dir)
        except zipfile.BadZipFile:
            raise HTTPException(400, "Invalid zip file")
        
        # Get folder tree (limit depth to 3)
        tree_items = []
        for item in extract_dir.rglob("*"):
            rel = item.relative_to(extract_dir)
            if len(rel.parts) <= 3:
                tree_items.append(str(rel))
            if len(tree_items) >= 50:
                break
        
        return {
            "folder_path": str(extract_dir),
            "tree": "\n".join(tree_items)
        }
        
    except HTTPException:
        raise
    except Exception as e:
        return {"success": False, "error": str(e)}


# =============================================================================
# WebSocket PTY Terminal
# =============================================================================
@app.websocket("/ws/terminal")
async def websocket_terminal(
    websocket: WebSocket, 
    token: str = Query(...),
    work_dir: str = Query(default=None),  # None = use server user's home
    session_id: str = Query(default="default")  # Session identifier for cwd tracking
):
    """WebSocket PTY - Full interactive terminal with custom work directory"""
    # Verify token
    if token != SECRET_TOKEN:
        await websocket.close(code=4001, reason="Unauthorized")
        return
    
    # Default to home folder of the user running the server
    if work_dir is None or work_dir == "" or work_dir == "~":
        work_dir = os.path.expanduser("~")
    
    # Validate and create work directory (expand ~ to home)
    work_path = Path(os.path.expanduser(work_dir)).resolve()
    work_path.mkdir(parents=True, exist_ok=True)
    
    await websocket.accept()
    
    # Create PTY
    master_fd, slave_fd = pty.openpty()
    
    # Set initial terminal size
    winsize = struct.pack("HHHH", 24, 80, 0, 0)
    fcntl.ioctl(slave_fd, termios.TIOCSWINSZ, winsize)
    
    # Environment for the shell
    env = os.environ.copy()
    env.update({
        "TERM": "xterm-256color",
        "COLORTERM": "truecolor",
        "LANG": "en_US.UTF-8",
        "LC_ALL": "en_US.UTF-8",
        "HOME": str(work_path),
        "PWD": str(work_path),
        "SHELL": "/bin/bash",
    })
    
    # Spawn bash as login shell
    pid = os.fork()
    if pid == 0:
        # Child process
        os.setsid()
        os.dup2(slave_fd, 0)
        os.dup2(slave_fd, 1)
        os.dup2(slave_fd, 2)
        os.close(master_fd)
        os.close(slave_fd)
        os.chdir(str(work_path))
        os.execvpe("bash", ["bash", "--login"], env)
    
    os.close(slave_fd)
    
    # Store session for cwd tracking
    pty_sessions[session_id] = {"pid": pid, "work_dir": str(work_path)}
    
    # Set non-blocking
    flags = fcntl.fcntl(master_fd, fcntl.F_GETFL)
    fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
    
    async def read_pty():
        """Read from PTY and send to WebSocket"""
        loop = asyncio.get_event_loop()
        while True:
            try:
                await asyncio.sleep(0.01)
                try:
                    data = os.read(master_fd, 4096)
                    if data:
                        await websocket.send_bytes(data)
                except BlockingIOError:
                    pass
                except OSError:
                    break
            except Exception:
                break
    
    async def write_pty():
        """Read from WebSocket and write to PTY"""
        while True:
            try:
                data = await websocket.receive()
                if data["type"] == "websocket.receive":
                    if "bytes" in data:
                        os.write(master_fd, data["bytes"])
                    elif "text" in data:
                        text = data["text"]
                        # Handle resize command
                        if text.startswith("RESIZE:"):
                            try:
                                _, cols, rows = text.split(":")
                                winsize = struct.pack("HHHH", int(rows), int(cols), 0, 0)
                                fcntl.ioctl(master_fd, termios.TIOCSWINSZ, winsize)
                            except:
                                pass
                        else:
                            os.write(master_fd, text.encode())
            except WebSocketDisconnect:
                break
            except Exception:
                break
    
    try:
        # Run both tasks
        read_task = asyncio.create_task(read_pty())
        write_task = asyncio.create_task(write_pty())
        
        done, pending = await asyncio.wait(
            [read_task, write_task],
            return_when=asyncio.FIRST_COMPLETED
        )
        
        for task in pending:
            task.cancel()
    finally:
        os.close(master_fd)
        pty_sessions.pop(session_id, None)
        try:
            os.kill(pid, signal.SIGTERM)
            os.waitpid(pid, 0)
        except:
            pass


@app.get("/cwd")
async def get_cwd(
    session_id: str = Query(default="default"),
    token: str = Depends(verify_token)
):
    """Get current working directory of a terminal session"""
    session = pty_sessions.get(session_id)
    if not session:
        raise HTTPException(404, "Session not found")
    
    pid = session["pid"]
    try:
        # Read the symlink /proc/{pid}/cwd to get actual cwd
        cwd = os.readlink(f"/proc/{pid}/cwd")
        return {"cwd": cwd, "session_id": session_id}
    except (OSError, FileNotFoundError):
        # Fallback to initial work_dir
        return {"cwd": session["work_dir"], "session_id": session_id}


def get_process_info(pid: int) -> Optional[dict]:
    """Read process info from /proc/{pid}"""
    try:
        proc_path = Path(f"/proc/{pid}")
        if not proc_path.exists():
            return None
        
        # Read stat file for basic info
        stat_content = (proc_path / "stat").read_text()
        stat_parts = stat_content.split()
        
        # Parse comm (process name) - it's in parentheses
        comm_start = stat_content.find('(')
        comm_end = stat_content.rfind(')')
        name = stat_content[comm_start+1:comm_end] if comm_start >= 0 and comm_end >= 0 else "unknown"
        
        # After comm, the rest of stat fields
        after_comm = stat_content[comm_end+2:].split()
        state = after_comm[0] if len(after_comm) > 0 else "?"
        ppid = int(after_comm[1]) if len(after_comm) > 1 else 0
        
        # Read cmdline for full command
        try:
            cmdline = (proc_path / "cmdline").read_text().replace('\x00', ' ').strip()
        except:
            cmdline = name
        
        # Read status for memory info
        rss_kb = 0
        try:
            status_content = (proc_path / "status").read_text()
            for line in status_content.split('\n'):
                if line.startswith("VmRSS:"):
                    rss_kb = int(line.split()[1])
                    break
        except:
            pass
        
        return {
            "pid": pid,
            "ppid": ppid,
            "name": name,
            "cmd": cmdline or name,
            "state": state,
            "memory_kb": rss_kb,
            "children": []
        }
    except (PermissionError, FileNotFoundError, OSError, ValueError, IndexError):
        return None


def build_process_tree(root_pid: Optional[int] = None) -> list:
    """Build hierarchical process tree"""
    processes = {}
    
    # Read all processes from /proc
    proc_path = Path("/proc")
    for entry in proc_path.iterdir():
        if entry.name.isdigit():
            pid = int(entry.name)
            info = get_process_info(pid)
            if info:
                processes[pid] = info
    
    # Build tree structure
    for pid, proc in processes.items():
        ppid = proc["ppid"]
        if ppid in processes and ppid != pid:
            processes[ppid]["children"].append(proc)
    
    # If root_pid specified, return subtree; otherwise return root processes
    if root_pid and root_pid in processes:
        return [processes[root_pid]]
    
    # Return only root processes (ppid=0 or ppid not in our list)
    roots = []
    for pid, proc in processes.items():
        if proc["ppid"] == 0 or proc["ppid"] not in processes:
            roots.append(proc)
    
    # Sort by PID
    roots.sort(key=lambda x: x["pid"])
    return roots


@app.get("/processes")
async def get_processes(
    session_id: Optional[str] = Query(default=None),
    token: str = Depends(verify_token)
):
    """Get process tree, optionally filtered by session's shell PID"""
    root_pid = None
    
    if session_id:
        session = pty_sessions.get(session_id)
        if session:
            root_pid = session["pid"]
    
    tree = build_process_tree(root_pid)
    return {"processes": tree, "session_id": session_id}


@app.post("/processes/kill")
async def kill_process_by_pid(
    pid: int = Query(...),
    sig: int = Query(default=15),  # SIGTERM by default
    token: str = Depends(verify_token)
):
    """Kill a process by PID"""
    try:
        os.kill(pid, sig)
        return {"status": "ok", "pid": pid, "signal": sig}
    except ProcessLookupError:
        raise HTTPException(404, f"Process {pid} not found")
    except PermissionError:
        raise HTTPException(403, f"Permission denied to kill process {pid}")
    except Exception as e:
        raise HTTPException(500, str(e))


# =============================================================================
# Static Files & Index
# =============================================================================
static_dir = Path(__file__).parent / "static"


@app.get("/health")
async def health(token: str = Depends(verify_token)):
    return {"status": "ok", "work_dir": str(WORK_DIR)}


# Mount static files (at the end so API routes take precedence)
if static_dir.exists():
    app.mount("/", StaticFiles(directory=str(static_dir), html=True), name="static")


if __name__ == "__main__":
    import uvicorn
    # Hugging Face Spaces uses port 7860 by default
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)