from fastapi import FastAPI, Request, Response
import redis
import aiohttp
import os
import logging

# Default configuration
CONFIG = {
    "REDIS_HOST": "127.0.0.1",
    "REDIS_PORT": "6379",
    "REDIS_DB": "2",
    "MINIO_ENDPOINT": "http://localhost:9000",
    "LOG_FILE": "/var/log/minio_quota.log"
}

# Load config from /etc/minio_quota.conf if available
CONFIG_FILE = "/etc/minio_quota.conf"
if os.path.exists(CONFIG_FILE):
    with open(CONFIG_FILE, "r") as f:
        for line in f:
            if "=" in line:
                key, value = line.strip().split("=", 1)
                CONFIG[key.strip()] = value.strip()

# Read configuration with priority:
REDIS_HOST = os.getenv("REDIS_HOST", CONFIG.get("REDIS_HOST"))
REDIS_PORT = os.getenv("REDIS_PORT", CONFIG.get("REDIS_PORT"))
REDIS_DB = int(os.getenv("REDIS_DB", CONFIG.get("REDIS_DB")))
MINIO_ENDPOINT = os.getenv("MINIO_ENDPOINT", CONFIG.get("MINIO_ENDPOINT"))
LOG_FILE = os.getenv("LOG_FILE", CONFIG.get("LOG_FILE"))

# Set up logging
handlers = [logging.StreamHandler()]  # Always log to stdout
if LOG_FILE and LOG_FILE.strip():  # Log to file only if LOG_FILE is set and not empty
    handlers.append(logging.FileHandler(LOG_FILE))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=handlers
)

# Initialize FastAPI and Redis client
app = FastAPI()
redis_client = redis.Redis(host=REDIS_HOST, port=int(REDIS_PORT), db=REDIS_DB)

@app.api_route("/{path:path}", methods=["GET", "PUT", "POST", "DELETE"])
async def proxy_request(path: str, request: Request):
    username = path.split("/")[0].split("-")[0]
    logging.info(f"Received {request.method} request for: {path}")

    if username:
        quota_exceeded = redis_client.get(f"quota_exceeded:{username}")
        if quota_exceeded:
            logging.warning(f"Quota exceeded for user: {username}, blocking request.")
            error_xml = f"""<?xml version="1.0" encoding="UTF-8"?>
<Error>
  <Code>QuotaExceeded</Code>
  <Message>User has exceeded storage quota.</Message>
  <Resource>/{path}</Resource>
  <RequestId>request-id-12345</RequestId>
</Error>"""
            return Response(content=error_xml, status_code=403, media_type="application/xml")

    query_string = request.url.query
    minio_url = f"{MINIO_ENDPOINT}/{path}"
    if query_string:
        minio_url += f"?{query_string}"
    headers = dict(request.headers)

    try:
        async with aiohttp.ClientSession() as session:
            async with session.request(
                method=request.method,
                url=minio_url,
                headers=headers,
                data=request.stream() if request.method in ["PUT", "POST"] else None
            ) as minio_response:
                
                # Try reading the response body
                try:
                    body = await minio_response.read()
                except aiohttp.ClientConnectionError:
                    body = b""  # If MinIO closes connection, return an empty response

                # Forward MinIO's response as-is
                return Response(
                    content=body,
                    status_code=minio_response.status,
                    headers=dict(minio_response.headers)
                )

    except aiohttp.ClientConnectionError:
        return Response(content=b"", status_code=200)  # Silently ignore MinIO disconnection

    except aiohttp.ClientResponseError as e:
        return Response(content=e.message, status_code=e.status)

    except aiohttp.ClientError:
        return Response(content="An error occurred while communicating with MinIO.", status_code=502)

    except Exception:
        return Response(content="Internal Server Error", status_code=500)