"""FastAPI proxy for the Z-Image generator frontend."""
import json
import os
import secrets
import time
import fcntl
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")))
USAGE_PATH = Path(os.getenv("USAGE_PATH", Path(__file__).with_name("usage.json")))
GALLERY_MAX_ITEMS = int(os.getenv("GALLERY_MAX_ITEMS", "500"))
WHITELIST_PATH = Path(os.getenv("WHITELIST_PATH", Path(__file__).with_name("whitelist.txt")))
ADMIN_ID = "86427531"

# --- Usage Store ---
class UsageStore:
    def __init__(self, path: Path):
        self.path = path
        if not self.path.exists():
            self._write({})

    def _read(self) -> dict:
        try:
            if not self.path.exists(): return {}
            with self.path.open("r", encoding="utf-8") as f:
                return json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            return {}

    def _write(self, data: dict):
        try:
            payload = json.dumps(data, ensure_ascii=False, indent=2)
            temp_path = self.path.with_suffix(".tmp_proxy")
            with temp_path.open("w", encoding="utf-8") as f:
                f.write(payload)
            temp_path.replace(self.path)
        except Exception as e:
            logger.error(f"Failed to write usage: {e}")

    def get_usage(self, user_id: str) -> dict:
        data = self._read()
        import datetime
        today = datetime.date.today().isoformat()
        user_data = data.get(user_id, {"daily_used": 0, "bonus_count": 0, "last_reset": today})
        
        if user_data.get("last_reset") != today:
            user_data["daily_used"] = 0
            user_data["last_reset"] = today
        return user_data

    def update_bonus(self, user_id: str, delta: int):
        data = self._read()
        import datetime
        today = datetime.date.today().isoformat()
        user_data = data.get(user_id, {"daily_used": 0, "bonus_count": 0, "last_reset": today})
        
        if user_data.get("last_reset") != today:
            user_data["daily_used"] = 0
            user_data["last_reset"] = today
            
        user_data["bonus_count"] = max(0, user_data.get("bonus_count", 0) + delta)
        data[user_id] = user_data
        self._write(data)

# --- 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")
    seed: Optional[int] = Field(default=None, ge=0)

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

    def delete_item(self, item_id: str) -> bool:
        with self.lock:
            data = self._read()
            items = data.get(self.item_key, [])
            initial_len = len(items)
            items = [i for i in items if i.get("id") != item_id]
            if len(items) < initial_len:
                data[self.item_key] = items
                self._write(data)
                return True
            return False

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)
usage_store = UsageStore(USAGE_PATH)

# --- App Setup ---
app = FastAPI(title="Z-Image Proxy", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_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")):
    is_liked_before = False
    items = image_store.list_items()
    target_item = next((i for i in items if i.get("id") == item_id), None)
    if target_item:
        is_liked_before = user_id in target_item.get("likedBy", [])
    
    updated_item = image_store.toggle_like(item_id, user_id)
    if updated_item:
        is_liked_after = user_id in updated_item.get("likedBy", [])
        if is_liked_after and not is_liked_before:
            usage_store.update_bonus(user_id, 1)
        elif not is_liked_after and is_liked_before:
            usage_store.update_bonus(user_id, -1)
        return updated_item
        
    updated_item = video_store.toggle_like(item_id, user_id)
    if updated_item:
        is_liked_after = user_id in updated_item.get("likedBy", [])
        if is_liked_after and not is_liked_before:
            usage_store.update_bonus(user_id, 1)
        elif not is_liked_after and is_liked_before:
            usage_store.update_bonus(user_id, -1)
        return updated_item
    raise HTTPException(status_code=404, detail="Item not found")

@app.get("/usage/{user_id}")
async def get_user_usage(user_id: str):
    try:
        usage = usage_store.get_usage(user_id)
        is_admin = user_id == ADMIN_ID
        remaining = (2 - usage["daily_used"]) + usage["bonus_count"] if not is_admin else 999999
        return {
            "daily_used": usage["daily_used"],
            "bonus_count": usage["bonus_count"],
            "base_limit": 2,
            "remaining": max(0, remaining),
            "is_admin": is_admin
        }
    except Exception as e:
        logger.error(f"Error getting usage for {user_id}: {e}")
        return {
            "daily_used": 0,
            "bonus_count": 0,
            "base_limit": 2,
            "remaining": 2,
            "is_admin": user_id == ADMIN_ID
        }

@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.delete("/gallery/videos/{item_id}")
async def delete_video(item_id: str, user_id: str = Query(..., alias="userId")):
    items = video_store.list_items()
    target_item = next((i for i in items if i.get("id") == item_id), None)
    
    if not target_item:
        raise HTTPException(status_code=404, detail="Video not found")
        
    ADMIN_ID = "86427531"
    
    if user_id != target_item.get("authorId") and user_id != ADMIN_ID:
        raise HTTPException(status_code=403, detail="Not authorized to delete this video")

    if video_store.delete_item(item_id):
        return {"status": "ok", "id": item_id}
    raise HTTPException(status_code=500, detail="Failed to delete video")

@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)