import logging
import asyncio
import os
import json
import hmac
from fastapi import FastAPI, HTTPException, Header, Depends, WebSocket, WebSocketDisconnect, Security, Query
from fastapi.security.api_key import APIKeyHeader, APIKey
from pydantic import BaseModel
from typing import List, Optional, Any, Dict
from ledgermind.core.api.memory import Memory
from sse_starlette.sse import EventSourceResponse
from starlette.concurrency import run_in_threadpool

logger = logging.getLogger("agent_memory_gateway")
app = FastAPI(title="Agent Memory REST & Real-time Gateway")

API_KEY_NAME = "X-API-Key"
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)

memory_instance: Optional[Memory] = None

class SearchRequest(BaseModel):
    query: str
    limit: int = 5
    offset: int = 0
    namespace: str = "default"
    mode: str = "balanced"

class RecordRequest(BaseModel):
    title: str
    target: str
    rationale: str
    consequences: Optional[List[str]] = []
    namespace: str = "default"

async def get_api_key(
    api_key_header: str = Security(api_key_header),
    api_key_query: Optional[str] = Query(None, alias="api_key")
):
    expected_key = os.environ.get("LEDGERMIND_API_KEY")
    if not expected_key:
        logger.error("LEDGERMIND_API_KEY is not set. Refusing access.")
        raise HTTPException(status_code=500, detail="Server authentication not configured")
    
    # Check both header and query param
    key = api_key_header or api_key_query
    if expected_key is None and key is None:
        return key
    if key is not None and expected_key is not None and hmac.compare_digest(key, expected_key):
        return key
    raise HTTPException(status_code=403, detail="Could not validate credentials")

def get_memory():
    if memory_instance is None:
        raise HTTPException(status_code=500, detail="Memory not initialized")
    return memory_instance

@app.post("/search", dependencies=[Depends(get_api_key)])
async def search(req: SearchRequest, mem: Memory = Depends(get_memory)):
    results = await run_in_threadpool(mem.search_decisions, req.query, limit=req.limit, mode=req.mode)
    return {"status": "success", "results": results}

@app.post("/record", dependencies=[Depends(get_api_key)])
async def record(req: RecordRequest, mem: Memory = Depends(get_memory)):
    try:
        res = mem.record_decision(
            title=req.title, target=req.target, 
            rationale=f"[via REST] {req.rationale}", 
            consequences=req.consequences
        )
        return {"status": "success", "id": res.metadata.get("file_id")}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/events", dependencies=[Depends(get_api_key)])
async def sse_events(mem: Memory = Depends(get_memory)):
    """Streaming endpoint for memory updates (SSE)."""
    async def event_generator():
        queue = asyncio.Queue()
        
        async def on_change(event_type, data):
            await queue.put({"event": event_type, "data": data})
            
        mem.events.subscribe(on_change)
        
        try:
            while True:
                change = await queue.get()
                yield {
                    "event": change["event"],
                    "data": json.dumps(change["data"])
                }
        finally:
            mem.events.unsubscribe(on_change)

    return EventSourceResponse(event_generator())

@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket, 
    mem: Memory = Depends(get_memory),
    api_key: Optional[str] = Query(None)
):
    """Bi-directional live updates via WebSockets."""
    expected_key = os.environ.get("LEDGERMIND_API_KEY")

    if not expected_key:
        logger.error("LEDGERMIND_API_KEY is not set. Refusing websocket connection.")
        await websocket.accept()
        await websocket.close(code=1008) # Policy Violation
        return

    if expected_key is not None and (api_key is None or not hmac.compare_digest(api_key, expected_key)):
        await websocket.accept()
        await websocket.close(code=4003) # Forbidden
        return

    await websocket.accept()
    
    async def on_change(event_type, data):
        try:
            await websocket.send_json({"event": event_type, "data": data})
        except Exception: 
            # Connection might be closed, subscription cleanup will happen in finally
            pass

    mem.events.subscribe(on_change)
    
    try:
        while True:
            # Handle incoming commands via WS if needed
            data = await websocket.receive_text()
            await websocket.send_json({"status": "received", "echo": data})
    except WebSocketDisconnect:
        logger.info("WebSocket client disconnected")
    finally:
        mem.events.unsubscribe(on_change)

@app.get("/health")
async def health():
    return {"status": "alive"}

def run_gateway(memory: Memory, host: str = "0.0.0.0", port: int = 8000): # nosec B104
    global memory_instance
    memory_instance = memory
    import uvicorn
    uvicorn.run(app, host=host, port=port)
