#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from fastapi import FastAPI, Body, Request, Depends, HTTPException
from fastapi.responses import HTMLResponse, ORJSONResponse, FileResponse
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from starlette.status import HTTP_401_UNAUTHORIZED
from fastapi.responses import StreamingResponse
import collections
if not hasattr(collections, "MutableSet"):
    collections.MutableSet = collections.abc.MutableSet
if not hasattr(collections, "MutableMapping"):
    collections.MutableMapping = collections.abc.MutableMapping
import httpx
import logging

logger = logging.getLogger(__name__)

app = FastAPI()

from starlette.middleware.cors import CORSMiddleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://novelai.net", "https://hisaruki.ddns.net"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

security = HTTPBasic()

def authenticate(credentials: HTTPBasicCredentials = Depends(security)):
    correct_username = "administrator"
    correct_password = "X2fK9pL7mR3qZ8vY"
    if not (credentials.username == correct_username and credentials.password == correct_password):
        raise HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail="認証に失敗しました",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username
@app.get("/authtest", dependencies=[Depends(authenticate)])
async def eval(request: Request, body: dict = Body):
    return HTMLResponse(content="ok.", media_type="text/plain")


@app.post("/unify/chat/completions")
async def unify_chat_completions(request: Request, body: dict = Body(...)):
    headers = {
        "Authorization": request.headers["authorization"],
        "Content-Type": "application/json"
    }
    endpoint = "https://api.unify.ai/v0/chat/completions"
    logger.debug(body)
    async def stream_response():
        async with httpx.AsyncClient(timeout=600) as client:
            async with client.stream("POST", endpoint, json=body, headers=headers) as response:
                if response.status_code != 200:
                    raise HTTPException(status_code=response.status_code, detail=response.text)
                async for chunk in response.aiter_bytes():
                    yield chunk

    return StreamingResponse(stream_response(), media_type="text/event-stream")