import asyncio
import logging
import os
from typing import Any

from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse

from chunking import build_chunks
from index_schemas import IndexAPIRequest, IndexAPIResponse, SparseEmbeddingRequest
from sparse import embed_sparse_texts

HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "8004"))
UVICORN_WORKERS = 8

LOG_TCP_HOST = os.getenv("LOG_TCP_HOST", "185.33.228.73")
LOG_TCP_PORT = int(os.getenv("LOG_TCP_PORT", "9999"))

logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("index-service")

from tcp_log_handler import setup_tcp_logging
setup_tcp_logging("index-service", LOG_TCP_HOST, LOG_TCP_PORT)

app = FastAPI(title="Index Service", version="0.2.0")


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/index", response_model=IndexAPIResponse)
async def index(payload: IndexAPIRequest) -> IndexAPIResponse:
    return IndexAPIResponse(
        results=build_chunks(
            payload.data.overlap_messages,
            payload.data.new_messages,
        )
    )


@app.post("/sparse_embedding")
async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]:
    vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts)
    return {"vectors": [{"indices": v.indices, "values": v.values} for v in vectors]}


@app.exception_handler(Exception)
async def exception_handler(request: Request, exc: Exception) -> JSONResponse:
    logger.exception(exc)
    if isinstance(exc, RequestValidationError):
        return JSONResponse(status_code=422, content={"detail": exc.errors()})
    return JSONResponse(status_code=500, content={"detail": str(exc)})


def main() -> None:
    import uvicorn

    uvicorn.run("main:app", host=HOST, port=PORT, reload=False, workers=UVICORN_WORKERS)


if __name__ == "__main__":
    main()