import os
import json
import base64
import time
import hmac
import hashlib
import secrets
from io import BytesIO
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, model_validator
from typing import Any
import torch
from diffusers import FluxPipeline
import bcrypt

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

print("Initializing Flux-Schnell on RTX 5060 Ti (sm_120)...")
hf_token = os.environ.get("HF_TOKEN")

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
    token=hf_token
)

# Offload execution stages sequentially to maintain 16GB headroom
pipe.enable_sequential_cpu_offload()
pipe.vae.enable_tiling()

# Auth config
USERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'users.json')
SESSION_SECRET = os.environ.get('SESSION_SECRET', secrets.token_hex(32))
SESSION_DURATION = 86400  # 24 hours


def load_users():
    with open(USERS_FILE) as f:
        return json.load(f)


def create_session(username: str) -> str:
    payload = f"{username}:{int(time.time())}"
    signature = hmac.new(
        SESSION_SECRET.encode(),
        payload.encode(),
        hashlib.sha256
    ).hexdigest()
    return f"{payload}.{signature}"


def verify_session(session: str) -> str | None:
    try:
        payload, sig = session.rsplit('.', 1)
        expected_sig = hmac.new(
            SESSION_SECRET.encode(),
            payload.encode(),
            hashlib.sha256
        ).hexdigest()
        if not hmac.compare_digest(sig, expected_sig):
            return None
        username, timestamp = payload.split(':', 1)
        if time.time() - int(timestamp) > SESSION_DURATION:
            return None
        users = load_users()
        if username not in users:
            return None
        return username
    except Exception:
        return None


class ImageRequest(BaseModel):
    prompt: str
    model: str = "flux-schnell"
    width: int = 1024
    height: int = 1024

    @model_validator(mode="before")
    @classmethod
    def validate_dimensions(cls, data: Any) -> Any:
        if isinstance(data, dict):
            width = data.get("width", 1024)
            height = data.get("height", 1024)
            if width % 32 != 0 or height % 32 != 0:
                raise ValueError("Width and height must be multiples of 32")
            if width > 2048 or height > 2048:
                raise ValueError("Maximum width and height is 2048")
        return data


@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    # Allow auth endpoints and static files through
    if request.url.path.startswith(("/auth/login", "/auth/logout", "/auth/check")):
        return await call_next(request)

    # Check session cookie for API requests
    session = request.cookies.get("session")
    username = verify_session(session) if session else None

    if not username:
        return JSONResponse(
            status_code=401,
            content={"detail": "Not authenticated"}
        )

    request.state.username = username
    return await call_next(request)


@app.post("/auth/login")
async def login(request: Request):
    data = await request.json()
    username = data.get('username', '')
    password = data.get('password', '')
    users = load_users()

    if username not in users:
        raise HTTPException(401, "Invalid credentials")

    if not bcrypt.checkpw(password.encode(), users[username].encode()):
        raise HTTPException(401, "Invalid credentials")

    session = create_session(username)
    response = JSONResponse(content={"username": username})
    response.set_cookie(
        key="session",
        value=session,
        httponly=True,
        samesite="strict",
        max_age=SESSION_DURATION,
        path="/",
    )
    return response


@app.post("/auth/logout")
async def logout():
    response = JSONResponse(content={"message": "Logged out"})
    response.set_cookie("session", "", expires=0, path="/")
    return response


@app.get("/auth/check")
async def check(request: Request, session: str | None = None):
    username = verify_session(session) if session else None
    if not username:
        raise HTTPException(401, "Not authenticated")
    return {"username": username}


@app.post("/v1/images/generations")
async def generate(request: Request):
    try:
        body = await request.json()
        req = ImageRequest(**body)
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

    try:
        result = pipe(
            req.prompt,
            guidance_scale=0.0,
            num_inference_steps=4,
            max_sequence_length=256,
            width=req.width,
            height=req.height,
        ).images[0]

        buffered = BytesIO()
        result.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        return {"data": [{"b64_json": img_str}]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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