from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from contextlib import asynccontextmanager
from tn3270_service import TN3270Session
import asyncio
import logging
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

session = None

@asynccontextmanager
async def lifespan(app):
    global session
    session = TN3270Session()
    try:
        session.connect()
        logger.info("TN3270 session ready")
    except Exception as e:
        logger.error(f"Failed to connect: {e}")
    yield
    if session:
        session.disconnect()

app = FastAPI(title="z/OS Web API", version="1.0.0", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

class CommandRequest(BaseModel):
    text: str

def check():
    if not session or not session.connected:
        raise HTTPException(503, "Not connected to z/OS")

@app.get("/health")
def health():
    return {"status": "ok", "connected": session.connected if session else False}

@app.get("/screen")
def get_screen():
    check()
    return {"screen": session.get_screen()}

@app.post("/command")
def send_command(req: CommandRequest):
    check()
    return {"screen": session.send_command(req.text)}

@app.post("/enter")
def send_enter():
    check()
    return {"screen": session.send_enter()}

@app.post("/pf/{key}")
def send_pf(key: int):
    check()
    if not 1 <= key <= 24:
        raise HTTPException(400, "PF key must be 1-24")
    return {"screen": session.send_pf(key)}

@app.post("/tab")
def send_tab():
    check()
    return {"screen": session.send_tab()}

class KeyRequest(BaseModel):
    key: str

class CharRequest(BaseModel):
    char: str

@app.post("/key")
def send_key(req: KeyRequest):
    check()
    # Map key names to s3270 commands
    key_map = {
        "Tab":       "Tab()",
        "BackTab":   "BackTab()",
        "Backspace": "BackSpace()",
        "Delete":    "Delete()",
        "Home":      "Home()",
        "End":       "FieldEnd()",
        "Insert":    "Insert()",
        "Clear":     "Clear()",
        "Attn":      "Attn()",
        "PA1":       "PA(1)",
        "PA2":       "PA(2)",
    }
    s3270_key = key_map.get(req.key, req.key)
    return {"screen": session.send_key(s3270_key)}

@app.post("/char")
def send_char(req: CharRequest):
    check()
    return {"screen": session.send_char(req.char)}

@app.get("/stream/screen")
async def stream_screen():
    async def event_generator():
        last_screen = ""
        while True:
            try:
                if session and session.connected:
                    if not session.lock.locked():
                        screen = session.get_screen()
                        if screen != last_screen:
                            data = json.dumps({"screen": screen})
                            yield f"data: {data}\n\n"
                            last_screen = screen
                await asyncio.sleep(0.5)
            except Exception as e:
                logger.error(f"Stream error: {e}")
                break
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        }
    )