from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import json
import asyncio
from concurrent.futures import ThreadPoolExecutor
from ..core.adapter import KogniTermAdapter

router = APIRouter()
adapter = KogniTermAdapter()
executor = ThreadPoolExecutor(max_workers=5)

@router.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            message_data = json.loads(data)
            user_message = message_data.get("message", "")
            
            if not user_message:
                continue

            # Process the message in a thread pool since it's a blocking generator
            loop = asyncio.get_event_loop()
            
            def run_invoke():
                # We need to add the user message to history first or pass it to invoke
                # The current terminal app adds it to history usually.
                # Actually LLMService.invoke takes 'history' but we need to manage it.
                
                # Simple implementation: just call the generator
                # We might want to yield chunks as they come
                responses = []
                for chunk in adapter.llm_service.invoke(history=adapter.agent_state.messages):
                    # In a more advanced implementation, we would send chunks directly here
                    # but this function runs in a thread, so we can't easily await websocket.send
                    responses.append(chunk)
                return responses

            # Improved approach: Send chunks via websocket from inside the thread?
            # Better: use a queue or just iterate in a thread and use loop.call_soon_threadsafe
            
            async def process_and_send():
                # This is a bit tricky with blocking generators. 
                # Let's try to iterate the generator in the threadpool and send chunks
                
                def iterate_gen():
                    for chunk in adapter.llm_service.invoke():
                        # We'll use this inside an async loop
                        yield chunk

                # For now, let's do a simple non-streaming return to verify things work
                # and then optimize for real streaming.
                
                # TO-DO: Implement real-time chunking via some mechanism
                pass

            # Temporary implementation: Run the whole thing and return
            # (We will improve this to be real streaming)
            
            # Let's actually try to do real streaming using a wrapper
            from kogniterm.core.history_manager import HumanMessage, AIMessage
            
            # Add human message to history
            adapter.agent_state.messages.append(HumanMessage(content=user_message))
            
            # Run the generator in a separate thread
            def generate():
                full_response = ""
                for chunk in adapter.llm_service.invoke(history=adapter.agent_state.messages):
                    # chunk can be AIMessage or str
                    content = chunk.content if hasattr(chunk, 'content') else str(chunk)
                    full_response += content
                    # Use call_soon_threadsafe to send through websocket
                    asyncio.run_coroutine_threadsafe(
                        websocket.send_json({"type": "chunk", "content": content}),
                        loop
                    )
                return full_response

            loop = asyncio.get_event_loop()
            await loop.run_in_executor(executor, generate)
            
            await websocket.send_json({"type": "done"})

    except WebSocketDisconnect:
        print("Client disconnected")
    except Exception as e:
        print(f"WebSocket error: {e}")
        await websocket.send_json({"type": "error", "content": str(e)})
