import os
import time
import random
import httpx
from typing import Dict
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse, Response

app = FastAPI()

BASE_URL = os.getenv("BASE_URL", "https://ollama.com")

# CLIENT API KEY (for auth)
MASTER_API_KEY = "ollama-proxy-free"

# BACKEND API KEYS (for Ollama)
API_KEYS = [
    "8ca25de51e554c099962b78b7ce0c9e9.Mp5dnqctR2zzq3g-NO_M-cjW",
    "dbd1d0c534964684a6d4678348ab8d30.ieDfmSYVnf0MmTjR-AIdNrW9",
    "37e81a6be4104fbfbfbe2ecf557a2c10.GoIbzpHebdM9ZcHarQ9A12Cp"
]

# KEY STATUS
key_status: Dict[str, Dict] = {
    k: {"fail": 0, "cooldown": 0, "status": "active"}
    for k in API_KEYS
}

def auth(req: Request):
    """Validate client API key"""
    client_key = req.headers.get("Authorization", "").replace("Bearer ", "")
    if client_key != MASTER_API_KEY:
        raise HTTPException(401, "Unauthorized")

def pick_key():
    now = time.time()
    valid = [k for k, v in key_status.items() 
             if v["status"] != "dead" and v["cooldown"] < now]
    
    if valid:
        return random.choice(valid)
    
    # Reset if all failed
    for k in key_status:
        key_status[k]["status"] = "active"
        key_status[k]["fail"] = 0
    
    return random.choice(API_KEYS)

def mark_fail(key: str):
    key_status[key]["fail"] += 1
    if key_status[key]["fail"] >= 3:
        key_status[key]["cooldown"] = time.time() + 60

def mark_ok(key: str):
    key_status[key]["fail"] = 0
    key_status[key]["status"] = "active"

@app.get("/")
def root():
    active = sum(1 for v in key_status.values() if v["status"] == "active")
    return {
        "status": "ok",
        "mode": "openai-proxy",
        "keys_total": len(API_KEYS),
        "keys_active": active
    }

@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy(req: Request, path: str):
    # Auth first
    auth(req)
    
    target_url = f"{BASE_URL}/v1/{path}"
    body = await req.body()
    
    headers = dict(req.headers)
    headers.pop("host", None)
    headers.pop("content-length", None)
    
    # Try each key
    tried = set()
    for _ in range(len(API_KEYS)):
        key = pick_key()
        if key in tried:
            continue
        tried.add(key)
        
        headers["Authorization"] = f"Bearer {key}"
        
        try:
            async with httpx.AsyncClient(timeout=60) as client:
                resp = await client.request(
                    method=req.method,
                    url=target_url,
                    headers=headers,
                    content=body,
                    params=req.query_params
                )
                
                if resp.status_code < 400:
                    mark_ok(key)
                    return Response(
                        content=resp.content,
                        status_code=resp.status_code,
                        headers=dict(resp.headers)
                    )
                
                mark_fail(key)
                continue
                
        except Exception as e:
            mark_fail(key)
            continue
    
    return JSONResponse({"error": "all keys failed"}, status_code=500)
# rebuild
# Testing build