"""FastAPI proxy for the Z-Image generator frontend."""
import json
import os
import secrets
import time
from pathlib import Path
from threading import Lock, RLock
from typing import List, Literal, Optional

import httpx
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, ConfigDict
import logging

# --- Constants ---
logger = logging.getLogger("uvicorn.error")
Z_IMAGE_BASE_URL = os.getenv("Z_IMAGE_BASE_URL", "http://106.120.52.146:39009").rstrip("/")
REQUEST_TIMEOUT_SECONDS = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "120"))
GALLERY_IMAGES_PATH = Path(os.getenv("GALLERY_IMAGES_PATH", Path(__file__).with_name("gallery_images.json")))
GALLERY_VIDEOS_PATH = Path(os.getenv("GALLERY_VIDEOS_PATH", Path(__file__).with_name("gallery_videos.json")))
GALLERY_MAX_ITEMS = int(os.getenv("GALLERY_MAX_ITEMS", "500"))
WHITELIST_PATH = Path(os.getenv("WHITELIST_PATH", Path(__file__).with_name("whitelist.txt")))

# --- Pydantic Models ---
# Define dependent models first to avoid forward reference issues.

class ImageGenerationPayload(BaseModel):
    model_config = ConfigDict(populate_by_name=True)
    prompt: str = Field(..., min_length=1, max_length=2048)
    height: int = Field(1024, ge=64, le=2048)
    width: int = Field(1024, ge=64, le=2048)
    num_inference_steps: int = Field(8, ge=1, le=200)
    guidance_scale: float = Field(0.0, ge=0.0, le=20.0)
    seed: Optional[int] = Field(default=None, ge=0)
    negative_prompt: Optional[str] = Field(default=None, max_length=2048)
    output_format: Literal["base64", "url"] = "base64"
    author_id: Optional[str] = Field(default=None, alias="authorId", min_length=1, max_length=64)

class GalleryItem(BaseModel):
    model_config = ConfigDict(populate_by_name=True)
    id: str
    prompt: str = Field(..., min_length=1, max_length=2048)
    url: str
    created_at: float = Field(default_factory=lambda: time.time() * 1000, alias="createdAt")
    author_id: Optional[str] = Field(default=None, alias="authorId")
    likes: int = 0
    is_mock: bool = Field(default=False, alias="isMock")
    liked_by: List[str] = Field(default_factory=list, alias="likedBy")

class GalleryImage(GalleryItem):
    height: int = Field(..., ge=64, le=2048)
    width: int = Field(..., ge=64, le=2048)
    num_inference_steps: int = Field(..., ge=1, le=200)
    guidance_scale: float = Field(..., ge=0.0, le=20.0)
    seed: int = Field(..., ge=0)
    negative_prompt: Optional[str] = None

class GalleryVideo(GalleryItem):
    generation_time: Optional[float] = Field(default=None, alias="generationTime")

class ImageGenerationResponse(BaseModel):
    image: Optional[str] = None
    url: Optional[str] = None
    time_taken: float = 0.0
    error: Optional[str] = None
    request_params: ImageGenerationPayload
    gallery_item: Optional[GalleryImage] = None # No forward ref needed now

# --- Data Stores ---

class WhitelistStore:
    def __init__(self, path: Path) -> None:
        self.path = path
        self.lock = RLock()
        if not self.path.exists(): self._write(["86427531"])
    def _read(self) -> List[str]:
        if not self.path.exists(): return []
        try:
            with self.path.open("r", encoding="utf-8") as f:
                return [line.strip() for line in f.read().splitlines() if line.strip()]
        except OSError: return []
    def _write(self, ids: List[str]) -> None:
        with self.lock:
            try:
                with self.path.open("w", encoding="utf-8") as f: f.write("\n".join(ids))
            except OSError as exc: print(f"[WARN] Failed to write whitelist: {exc}")
    def is_allowed(self, user_id: str) -> bool: return user_id in self._read()
    def add_users(self, user_ids: List[str]) -> None:
        with self.lock:
            current = set(self._read()); current.update(user_ids); self._write(sorted(list(current)))
    def remove_user(self, user_id: str) -> None:
        with self.lock:
            current = self._read()
            if user_id in current: self._write([uid for uid in current if uid != user_id])
    def get_all(self) -> List[str]: return self._read()

class JsonStore:
    """Generic JSON file backed store for a list of items."""
    def __init__(self, path: Path, item_key: str, max_items: int = 500) -> None:
        self.path = path
        self.item_key = item_key
        self.max_items = max_items
        self.lock = Lock()
        try:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            if not self.path.exists(): self._write({self.item_key: []})
        except OSError as exc: print(f"[WARN] JSON store at {path} disabled: {exc}")
    def _read(self) -> dict:
        try:
            with self.path.open("r", encoding="utf-8") as file: return json.load(file)
        except (FileNotFoundError, json.JSONDecodeError): return {self.item_key: []}
    def _write(self, data: dict) -> None:
        payload = json.dumps(data, ensure_ascii=False, indent=2)
        temp_path = self.path.with_suffix(".tmp")
        try:
            with temp_path.open("w", encoding="utf-8") as file: file.write(payload)
            temp_path.replace(self.path)
        except OSError:
            with self.path.open("w", encoding="utf-8") as file: file.write(payload)
    def list_items(self) -> List[dict]:
        with self.lock: return self._read().get(self.item_key, [])
    def add_item(self, item: BaseModel) -> dict:
        payload = item.model_dump(by_alias=True)
        with self.lock:
            data = self._read()
            items = data.get(self.item_key, [])
            items.insert(0, payload)
            data[self.item_key] = items[:self.max_items]
            self._write(data)
        return payload
    def toggle_like(self, item_id: str, user_id: str) -> Optional[dict]:
        with self.lock:
            data = self._read()
            items = data.get(self.item_key, [])
            target_item = next((i for i in items if i.get("id") == item_id), None)
            if not target_item: return None
            liked_by = target_item.get("likedBy", [])
            if not isinstance(liked_by, list): liked_by = []
            if user_id in liked_by:
                liked_by.remove(user_id)
                target_item["likes"] = max(0, target_item.get("likes", 0) - 1)
            else:
                liked_by.append(user_id)
                target_item["likes"] = target_item.get("likes", 0) + 1
            target_item["likedBy"] = liked_by
            self._write(data)
            return target_item

image_store = JsonStore(GALLERY_IMAGES_PATH, item_key="images", max_items=GALLERY_MAX_ITEMS)
video_store = JsonStore(GALLERY_VIDEOS_PATH, item_key="videos", max_items=GALLERY_MAX_ITEMS)
whitelist_store = WhitelistStore(WHITELIST_PATH)

# --- App Setup ---
app = FastAPI(title="Z-Image Proxy", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://106.120.52.146:37001"],  # Explicitly allow the frontend origin
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
@app.on_event("startup")
async def startup(): app.state.http = httpx.AsyncClient(timeout=httpx.Timeout(REQUEST_TIMEOUT_SECONDS, connect=5.0))
@app.on_event("shutdown")
async def shutdown(): await app.state.http.aclose()
@app.get("/health")
async def health(): return {"status": "ok"}

# --- Endpoints ---
@app.post("/auth/login")
async def login(user_id: str = Query(..., alias="userId")):
    if whitelist_store.is_allowed(user_id): return {"status": "ok", "userId": user_id}
    raise HTTPException(status_code=403, detail="User not whitelisted")

@app.post("/likes/{item_id}")
async def toggle_like(item_id: str, user_id: str = Query(..., alias="userId")):
    updated_item = image_store.toggle_like(item_id, user_id)
    if updated_item: return updated_item
    updated_item = video_store.toggle_like(item_id, user_id)
    if updated_item: return updated_item
    raise HTTPException(status_code=404, detail="Item not found")

@app.get("/gallery/images")
async def gallery_images(limit: int = Query(200, ge=1, le=1000), author_id: Optional[str] = Query(None, alias="authorId")):
    items = image_store.list_items()
    if author_id: items = [item for item in items if item.get("authorId") == author_id]
    return {"images": items[:limit]}

@app.get("/gallery/videos")
async def gallery_videos(limit: int = Query(200, ge=1, le=1000), author_id: Optional[str] = Query(None, alias="authorId")):
    items = video_store.list_items()
    if author_id: items = [item for item in items if item.get("authorId") == author_id]
    return {"videos": items[:limit]}

@app.post("/gallery/videos")
async def add_video(video: GalleryVideo):
    try:
        return video_store.add_item(video)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to store video metadata: {exc}")

@app.post("/generate", response_model=ImageGenerationResponse)
async def generate_image(payload: ImageGenerationPayload):
    request_params_data = payload.model_dump(); body = {k: v for k, v in request_params_data.items() if v is not None and k != "author_id"}
    if "seed" not in body: body["seed"] = secrets.randbelow(1_000_000_000)
    request_params_data["seed"] = body["seed"]; request_params = ImageGenerationPayload(**request_params_data)
    url = f"{Z_IMAGE_BASE_URL}/generate"
    try:
        resp = await app.state.http.post(url, json=body)
        resp.raise_for_status()
        data = resp.json()
        image_url = data.get("url") or f"data:image/png;base64,{data.get('image')}"
        stored_item_data = {
            "id": data.get("id") or secrets.token_hex(16),
            "prompt": payload.prompt,
            "width": payload.width,
            "height": payload.height,
            "num_inference_steps": payload.num_inference_steps,
            "guidance_scale": payload.guidance_scale,
            "seed": request_params.seed,
            "url": image_url,
            "author_id": payload.author_id,
            "negative_prompt": payload.negative_prompt,
        }
        stored = image_store.add_item(GalleryImage(**stored_item_data))
        return ImageGenerationResponse(image=data.get("image"), url=data.get("url"), time_taken=data.get("time_taken", 0.0), request_params=request_params, gallery_item=GalleryImage.model_validate(stored))
    except httpx.RequestError as exc: raise HTTPException(status_code=502, detail=f"Z-Image service unreachable: {exc}")
    except Exception as exc: raise HTTPException(status_code=500, detail=f"An error occurred: {exc}")

@app.get("/admin/whitelist")
async def get_whitelist() -> dict:
    return {"whitelist": whitelist_store.get_all()}
@app.post("/admin/whitelist")
async def add_whitelist(user_ids: List[str]) -> dict:
    whitelist_store.add_users(user_ids)
    return {"status": "ok", "whitelist": whitelist_store.get_all()}
@app.delete("/admin/whitelist/{user_id}")
async def remove_whitelist(user_id: str) -> dict:
    whitelist_store.remove_user(user_id)
    return {"status": "ok", "whitelist": whitelist_store.get_all()}

# Redirect old /gallery to /gallery/images for backward compatibility
@app.get("/gallery")
async def gallery(limit: int = Query(200, ge=1, le=1000), author_id: Optional[str] = Query(None, alias="authorId")):
    return await gallery_images(limit=limit, author_id=author_id)