import os
import pty
import asyncio
import struct
import fcntl
import termios
import json
import subprocess
import uuid
from pathlib import Path

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

PERSISTENT_DIR = Path("/persistent")
PKGS_FILE = PERSISTENT_DIR / "installed_packages.txt"
PERSISTENT_DIR.mkdir(parents=True, exist_ok=True)

# In-memory session storage - ptys stay alive even after WebSocket disconnects
sessions = {}

@app.get("/")
async def index():
    return FileResponse("static/index.html")

@app.get("/api/pass")
async def get_pass():
    return {"pass": os.environ.get("TERMINAL_PASS", "")}

@app.get("/api/sessions")
async def list_sessions():
    alive = []
    for sid, s in sessions.items():
        if s.get("proc") and s["proc"].returncode is None:
            alive.append({"id": sid, "alive": True})
    return JSONResponse(alive)

@app.websocket("/ws")
async def terminal(websocket: WebSocket):
    await websocket.accept()

    session_id = None
    read_task = None
    master_fd = None
    proc = None

    try:
        while True:
            message = await websocket.receive_text()
            msg = json.loads(message)

            if msg.get("type") == "attach":
                session_id = msg.get("session_id")

                # Reuse existing session if alive
                if session_id and session_id in sessions:
                    s = sessions[session_id]
                    if s.get("proc") and s["proc"].returncode is None:
                        master_fd = s["master_fd"]
                        proc = s["proc"]
                        await websocket.send_text(json.dumps({
                            "type": "session_info",
                            "session_id": session_id
                        }))
                        # Send current terminal state
                        try:
                            output = s.get("last_output", "")
                            if output:
                                await websocket.send_text(json.dumps({
                                    "type": "output",
                                    "data": output[-4000:]
                                }))
                        except:
                            pass
                        if read_task:
                            read_task.cancel()
                        read_task = asyncio.create_task(read_pty(websocket, master_fd, session_id))
                        continue

                # Create new session
                session_id = session_id or str(uuid.uuid4())[:8]
                master_fd, slave_fd = pty.openpty()

                env = {
                    **os.environ,
                    "TERM": "xterm-256color",
                    "HOME": "/persistent",
                    "USER": "ubuntu",
                    "SHELL": "/bin/bash",
                }

                proc = await asyncio.create_subprocess_exec(
                    "/bin/bash", "--login",
                    stdin=slave_fd, stdout=slave_fd, stderr=slave_fd,
                    env=env
                )
                os.close(slave_fd)

                sessions[session_id] = {
                    "master_fd": master_fd,
                    "proc": proc,
                    "last_output": ""
                }

                await websocket.send_text(json.dumps({
                    "type": "session_info",
                    "session_id": session_id
                }))

                if read_task:
                    read_task.cancel()
                read_task = asyncio.create_task(read_pty(websocket, master_fd, session_id))

            elif msg.get("type") == "input":
                if master_fd is not None:
                    try:
                        os.write(master_fd, msg["data"].encode("utf-8"))
                    except OSError:
                        pass

            elif msg.get("type") == "resize":
                if master_fd is not None:
                    try:
                        fcntl.ioctl(
                            master_fd, termios.TIOCSWINSZ,
                            struct.pack("HHHH", msg.get("rows", 24), msg.get("cols", 80), 0, 0)
                        )
                    except OSError:
                        pass

    except (WebSocketDisconnect, Exception):
        pass
    finally:
        if read_task:
            read_task.cancel()
        # DON'T kill the process - keep it alive for reconnection
        # Just clean up the WebSocket reference

async def read_pty(websocket: WebSocket, master_fd: int, session_id: str):
    loop = asyncio.get_event_loop()
    try:
        while True:
            data = await loop.run_in_executor(None, lambda: os.read(master_fd, 4096))
            if data:
                try:
                    text = data.decode("utf-8")
                except UnicodeDecodeError:
                    text = data.decode("utf-8", errors="ignore")

                # Cache last output for reconnect
                if session_id in sessions:
                    sessions[session_id]["last_output"] = sessions[session_id].get("last_output", "") + text
                    # Keep only last 8000 chars
                    if len(sessions[session_id]["last_output"]) > 8000:
                        sessions[session_id]["last_output"] = sessions[session_id]["last_output"][-4000:]

                try:
                    await websocket.send_text(json.dumps({"type": "output", "data": text}))
                except:
                    break
    except (OSError, asyncio.CancelledError):
        pass