from datetime import datetime
from json import dumps, loads
from re import sub
from urllib.parse import urlparse, urlunparse

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse, HTMLResponse
from httpx import AsyncClient, Limits, Timeout

app = FastAPI(title='PROXI-API')

MODELS = [
    'deepseek-ai/DeepSeek-R1',
    'deepseek-ai/DeepSeek-V3',
    'deepseek-ai/deepseek-llm-67b-chat',
    'databricks/dbrx-instruct',
    'Qwen/QwQ-32B-Preview',
    'NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO'
]

MODELS_OBJ = {
    'object': 'list',
    'data': [
        {
            'id': model,
            'object': 'model',
            'created': int(datetime.now().timestamp()),
            'owned_by': 'blackbox.ai'
        }
        for model in MODELS
    ]
}

BLOCKED_HEADERS = [
    'x-forwarded-for',
    'x-real-ip',
    'proxy-authorization'
]

HTTPX_CLIENT_KWARGS = dict(
    timeout=Timeout(
        connect=15,
        read=60,
        write=30,
        pool=30
    ),
    limits=Limits(
        max_keepalive_connections=10,
        max_connections=100
    ),
    follow_redirects=True,
    http2=True,
    verify=False
)


def normalize_headers(request: Request, original_url: str) -> dict:
    headers = {
        key: value for (key, value) in request.headers.items()
        if key.lower() not in BLOCKED_HEADERS and key.lower() != 'host'
    }
    parsed_url = urlparse(original_url)
    origin = urlunparse((parsed_url.scheme, parsed_url.netloc, '', '', '', ''))
    normalized = {key.lower(): value for key, value in headers.items()}
    normalized['referer'] = original_url
    normalized['origin'] = origin
    normalized['accept-encoding'] = 'deflate'
    return normalized


def decode(chunk: bytes) -> str:
    try:
        content = chunk.decode('utf-8')
    except UnicodeDecodeError:
        try:
            content = chunk.decode('latin-1')
        except UnicodeDecodeError:
            try:
                content = chunk.decode('iso-8859-1')
            except UnicodeDecodeError:
                content = chunk.decode(errors='replace')
    return content


def format_chunk(chunk: bytes, model: str) -> bytes | str | dict:
    chunk_id = 'chatcmpl-AQ8Lzxlg8eSCB1lgVmboiXwZiexqE'
    timestamp = int(datetime.now().timestamp())
    
    data = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "created": timestamp,
        "model": model,
        "system_fingerprint": "fp_67802d9a6d",
        "choices": [
            {
                "index": 0,
                "delta": {
                    "content": decode(chunk)
                },
                "finish_reason": None
            }
        ]
    }
    json_data = dumps(data, ensure_ascii=False)
    str_data = f'data: {sub(r'\\\\([ntr])', r'\\\1', json_data)}\n\n'
    return str_data


async def generate(request: Request, url: str, headers: dict, body: bytes):
    body_str = body.decode('utf-8')
    body_obj: dict = loads(body_str)
    model = body_obj.get('model')
    if 'max_tokens' not in body_obj:
        body_obj['max_tokens'] = '32000'
        body = dumps(body_obj, ensure_ascii=False).encode()
        headers = dict(headers)
        headers['content-length'] = str(len(body))
        headers['content-type'] = 'application/json'
    async with AsyncClient(**HTTPX_CLIENT_KWARGS) as stream_client:
        async with stream_client.stream(method=request.method, url=url, headers=headers, content=body, cookies=request.cookies) as stream:
            async for chunk in stream.aiter_raw():
                yield format_chunk(chunk, model)
            yield 'data: [DONE]\n\n'.encode()


@app.post('/api/chat/completions')
@app.post('/api/v1/chat/completions')
async def proxy(request: Request):
    try:
        url = 'https://api.blackbox.ai/api/chat'
        headers = normalize_headers(request, url)
        body = await request.body()
        async with AsyncClient(**HTTPX_CLIENT_KWARGS) as headers_client:
            async with headers_client.stream(method=request.method, url=url, headers=headers, content=body, cookies=request.cookies) as response:
                excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
                filtered_headers = {
                    name: value for (name, value) in response.headers.items()
                    if name.lower() not in excluded_headers
                }
                return StreamingResponse(generate(request, url, headers, body), status_code=response.status_code, headers=filtered_headers)

    except Exception as exc:
        return JSONResponse({'error': response.status_code, 'body': f'{exc}'}, status_code=500)


@app.get('/')
@app.get('/api')
@app.get('/api/v1')
async def root():
    return HTMLResponse('ну пролапс, ну и что', status_code=200)


@app.get('/api/models')
@app.get('/api/v1/models')
async def models():
    return JSONResponse(MODELS_OBJ, status_code=200, media_type='application/json')


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