from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Dict, List
import uuid
from datetime import datetime
from models import User, Stream, StreamCreate, AuthCallback, ChatMessage
from auth import oidc_client, get_current_user, create_access_token
from config import settings

app = FastAPI(title="Video Streaming Platform")

app.add_middleware(
    CORSMiddleware,
    allow_origins=[settings.frontend_url],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# In-memory storage
streams: Dict[str, Stream] = {}
users: Dict[str, User] = {}
active_connections: Dict[str, List[WebSocket]] = {}

@app.get("/")
async def root():
    return {"message": "Video Streaming API"}

@app.get("/auth/login")
async def login():
    discovery = await oidc_client.get_discovery_document()
    auth_endpoint = discovery["authorization_endpoint"]
    
    auth_url = (
        f"{auth_endpoint}?"
        f"client_id={settings.oidc_client_id}&"
        f"redirect_uri={settings.oidc_redirect_uri}&"
        f"response_type=code&"
        f"scope=openid email profile"
    )
    
    return {"auth_url": auth_url}

@app.post("/auth/callback")
async def auth_callback(callback: AuthCallback):
    try:
        token_response = await oidc_client.exchange_code(callback.code)
        access_token = token_response.get("access_token")
        
        if not access_token:
            raise HTTPException(status_code=400, detail="Failed to get access token")
        
        user_info = await oidc_client.get_user_info(access_token)
        
        user = User(
            id=user_info.get("sub"),
            email=user_info.get("email"),
            name=user_info.get("name", user_info.get("email")),
            picture=user_info.get("picture")
        )
        
        users[user.id] = user
        
        jwt_token = create_access_token({
            "sub": user.id,
            "email": user.email,
            "name": user.name
        })
        
        return {"access_token": jwt_token, "user": user}
    
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/streams")
async def get_streams():
    return list(streams.values())

@app.get("/streams/{stream_id}")
async def get_stream(stream_id: str):
    if stream_id not in streams:
        raise HTTPException(status_code=404, detail="Stream not found")
    return streams[stream_id]

@app.post("/streams")
async def create_stream(
    stream_data: StreamCreate,
    current_user: dict = Depends(get_current_user)
):
    stream_id = str(uuid.uuid4())
    stream = Stream(
        id=stream_id,
        user_id=current_user["sub"],
        title=stream_data.title,
        description=stream_data.description,
        is_live=True,
        created_at=datetime.now()
    )
    
    streams[stream_id] = stream
    active_connections[stream_id] = []
    
    return stream

@app.delete("/streams/{stream_id}")
async def end_stream(
    stream_id: str,
    current_user: dict = Depends(get_current_user)
):
    if stream_id not in streams:
        raise HTTPException(status_code=404, detail="Stream not found")
    
    stream = streams[stream_id]
    if stream.user_id != current_user["sub"]:
        raise HTTPException(status_code=403, detail="Not authorized")
    
    stream.is_live = False
    
    for connection in active_connections.get(stream_id, []):
        await connection.close()
    
    active_connections[stream_id] = []
    
    return {"message": "Stream ended"}

@app.websocket("/ws/stream/{stream_id}")
async def websocket_stream(websocket: WebSocket, stream_id: str):
    await websocket.accept()
    
    if stream_id not in streams:
        await websocket.close(code=1008)
        return
    
    if stream_id not in active_connections:
        active_connections[stream_id] = []
    
    active_connections[stream_id].append(websocket)
    streams[stream_id].viewer_count = len(active_connections[stream_id])
    
    try:
        while True:
            data = await websocket.receive_json()
            
            for connection in active_connections[stream_id]:
                if connection != websocket:
                    await connection.send_json(data)
    
    except WebSocketDisconnect:
        active_connections[stream_id].remove(websocket)
        streams[stream_id].viewer_count = len(active_connections[stream_id])

@app.get("/me")
async def get_me(current_user: dict = Depends(get_current_user)):
    user_id = current_user["sub"]
    if user_id in users:
        return users[user_id]
    return current_user

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)