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

import httpx
from fastapi import FastAPI, HTTPException, Query, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, ConfigDict
import logging
from PIL import Image
import io
import sys

# Add parent directory to path to import middleware
sys.path.append(str(Path(__file__).parent.parent))
try:
    from middleware import IPFilterMiddleware
except ImportError:
    # Fallback/Dummy if not found (should be found)
    IPFilterMiddleware = None


# --- 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(__file__).with_name("gallery_images.json")
GALLERY_VIDEOS_PATH = Path(__file__).with_name("gallery_videos.json")
USAGE_PATH = Path(__file__).with_name("usage.json")
GALLERY_MAX_ITEMS = int(os.getenv("GALLERY_MAX_ITEMS", "500"))
WHITELIST_PATH = Path(__file__).with_name("whitelist.txt")
ADMIN_ID = "86427531"

# --- Add this line to fix the port ---
TURBO_DIFFUSION_LOCAL_URL = os.getenv("TURBO_DIFFUSION_LOCAL_URL", "http://localhost:8000").rstrip("/")


# Load dynamic limits from config.js
CONFIG_JS_PATH = Path(__file__).parent.parent / "public" / "config.js"

def load_limits_from_config() -> dict:
    defaults = {"VIDEO_GENERATION_LIMIT": 1, "LIKES_FOR_REWARD": 5}
    try:
        if not CONFIG_JS_PATH.exists():
            return defaults
        
        content = CONFIG_JS_PATH.read_text(encoding="utf-8")
        
        # Simple regex to extract values from JS object
        # Looking for: VIDEO_GENERATION_LIMIT: 1
        limit_match = re.search(r'VIDEO_GENERATION_LIMIT\s*:\s*(\d+)', content)
        reward_match = re.search(r'LIKES_FOR_REWARD\s*:\s*(\d+)', content)
        
        if limit_match:
            defaults["VIDEO_GENERATION_LIMIT"] = int(limit_match.group(1))
        if reward_match:
            defaults["LIKES_FOR_REWARD"] = int(reward_match.group(1))
            
        return defaults
    except Exception as e:
        logger.error(f"Failed to load config.js: {e}")
        return defaults

LIMITS = load_limits_from_config()

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

    def _read(self) -> dict:
        lock_file = open(self.lock_path, "w")
        try:
            fcntl.flock(lock_file, fcntl.LOCK_EX)
            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 {}
        finally:
            fcntl.flock(lock_file, fcntl.LOCK_UN)
            lock_file.close()

    def _write(self, data: dict):
        lock_file = open(self.lock_path, "w")
        try:
            fcntl.flock(lock_file, fcntl.LOCK_EX)
            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}")
        finally:
            fcntl.flock(lock_file, fcntl.LOCK_UN)
            lock_file.close()

    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 increment_used(self, user_id: str):
        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
        
        if user_data["daily_used"] < LIMITS["VIDEO_GENERATION_LIMIT"]:
            user_data["daily_used"] += 1
        else:
            user_data["bonus_count"] = max(0, user_data.get("bonus_count", 0) - 1)
            
        data[user_id] = user_data
        self._write(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)
    width: int = Field(1024, ge=64, le=2048)
    height: int = Field(1024, ge=64, le=2048)
    thumbnail: Optional[str] = None

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 = []
            
            # --- New Reward Logic ---
            # 1. Check current likes BEFORE change
            current_likes_count = target_item.get("likes", 0)
            author_id = target_item.get("authorId")
            
            is_liked_after = False
            
            if user_id in liked_by:
                # UNLIKE
                liked_by.remove(user_id)
                new_likes_count = max(0, current_likes_count - 1)
                is_liked_after = False
            else:
                # LIKE
                liked_by.append(user_id)
                new_likes_count = current_likes_count + 1
                is_liked_after = True
            
            target_item["likes"] = new_likes_count
            target_item["likedBy"] = liked_by
            self._write(data)
            
            # Reward Check: Only reward author when crossing threshold (e.g. 5, 10, 15...)
            # We check if the NEW count is a multiple of LIKES_FOR_REWARD and we just increased it.
            # (Simple version: Every N likes = 1 generation credit)
            if author_id and author_id != "OFFICIAL" and author_id != ADMIN_ID:
                limit = LIMITS["LIKES_FOR_REWARD"]
                # Only reward on LIKE action, not unlike
                if is_liked_after:
                     # Check if we just hit a multiple of the limit (5, 10, 15...)
                     if new_likes_count > 0 and new_likes_count % limit == 0:
                         logger.info(f"User {author_id} reached {new_likes_count} likes! Adding bonus.")
                         usage_store.update_bonus(author_id, 1)

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

# IP Filter Middleware (Add BEFORE CORS to block early)
if IPFilterMiddleware:
    app.add_middleware(IPFilterMiddleware)

from fastapi.staticfiles import StaticFiles
# Mount public directory to serve thumbnails and frontend config
app.mount("/thumbnails", StaticFiles(directory=str(Path(__file__).parent.parent / "public" / "thumbnails")), name="thumbnails")

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

async def startup(): 
    # Reload limits on startup to ensure fresh config
    global LIMITS
    LIMITS = load_limits_from_config()
    logger.info(f"Loaded limits: {LIMITS}")
    app.state.http = httpx.AsyncClient(timeout=httpx.Timeout(REQUEST_TIMEOUT_SECONDS, connect=5.0), trust_env=False)

@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")):
    # Try images first
    updated_item = image_store.toggle_like(item_id, user_id)
    if updated_item: return updated_item
        
    # Then videos
    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("/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
        
        limit = LIMITS["VIDEO_GENERATION_LIMIT"]
        
        # Logic: base_limit - daily_used + bonus
        remaining = (limit - usage["daily_used"]) + usage["bonus_count"]
        if is_admin: remaining = 999999
        
        return {
            "daily_used": usage["daily_used"],
            "bonus_count": usage["bonus_count"],
            "base_limit": limit,
            "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": LIMITS["VIDEO_GENERATION_LIMIT"],
            "remaining": LIMITS["VIDEO_GENERATION_LIMIT"],
            "is_admin": user_id == ADMIN_ID
        }

@app.post("/submit-video-job/")
async def submit_video_job_proxy(
    prompt: str = Form(...),
    image: UploadFile = File(...),
    author_id: str = Form(...),
    num_steps: Optional[int] = Form(4),
    seed: Optional[int] = Form(0)
):
    
    # 1. Check Usage
    if author_id != ADMIN_ID:
        usage = usage_store.get_usage(author_id)
        limit = LIMITS["VIDEO_GENERATION_LIMIT"]
        allowed = (limit + usage.get("bonus_count", 0))
        if usage["daily_used"] >= allowed:
            logger.warning(f"User {author_id} limit reached via proxy")
            raise HTTPException(status_code=403, detail="今日生成次数已用完。点赞灵感图库的图片可增加次数！")

    # 2. Forward to TurboDiffusion
    url = f"{TURBO_DIFFUSION_LOCAL_URL}/submit-job/"
    
    # Prepare files and data for forwarding
    file_content = await image.read()
    
    # --- Save Thumbnail ---
    # Generate a lightweight thumbnail (max 400px width/height, ~20KB)
    thumbnail_filename = f"thumb_{secrets.token_hex(8)}.jpg" # Always use jpg for efficiency
    
    # Updated Path: Save to OSS directory
    OSS_THUMBNAIL_DIR = Path("/home/inspur/work_space/gen_img_video/TurboDiffusion-Space/ASSERT/艺云-DESIGN/thumbnails")
    thumbnail_path = OSS_THUMBNAIL_DIR / thumbnail_filename
    
    # Updated URL: Construct from environment variables (PUBLIC_IP and PUBLIC_OSS_PORT)
    public_ip = os.getenv("PUBLIC_IP", "106.120.52.146")
    oss_port = os.getenv("PUBLIC_OSS_PORT", "34000")
    oss_base_url = f"http://{public_ip}:{oss_port}"
    thumbnail_url_path = f"{oss_base_url}/thumbnails/{thumbnail_filename}"
    
    try:
        thumbnail_path.parent.mkdir(parents=True, exist_ok=True)
        # Use PIL to resize and compress
        with Image.open(io.BytesIO(file_content)) as img:
            # Convert to RGB to handle PNG/RGBA correctly
            if img.mode in ("RGBA", "P"):
                img = img.convert("RGB")
            
            # Resize while maintaining aspect ratio, max 400px
            img.thumbnail((400, 400))
            
            # Save optimized JPEG
            img.save(thumbnail_path, "JPEG", quality=70, optimize=True)
            logger.info(f"Saved optimized thumbnail to {thumbnail_path}")
            
    except Exception as e:
        logger.error(f"Failed to generate thumbnail: {e}")
        # Fallback: try to write original file if PIL fails, but rename extension if needed
        try:
             with open(thumbnail_path, "wb") as f:
                f.write(file_content)
        except Exception as e2:
             logger.error(f"Failed to save fallback thumbnail: {e2}")
             thumbnail_url_path = None

    files = {
        'image': (image.filename, file_content, image.content_type)
    }
    
    # Use the actual author_id now that we've disabled the check on the inference server
    data = {
        'prompt': prompt,
        'author_id': author_id, 
        'num_steps': str(num_steps),
        'seed': str(seed)
    }
    
    try:
        # Use a separate client or the app state client (but need to handle multipart)
        # Using separate httpx call for simplicity with files
        async with httpx.AsyncClient(timeout=30.0, trust_env=False) as client:
            resp = await client.post(url, data=data, files=files)
            
        if resp.status_code != 202 and resp.status_code != 200:
             raise HTTPException(status_code=resp.status_code, detail=f"Inference service error: {resp.text}")
             
        result = resp.json()
        
        # Inject thumbnail URL into the response so frontend can use it
        if thumbnail_url_path:
            logger.info(f"Injecting thumbnail URL: {thumbnail_url_path}")
            result["thumbnail"] = thumbnail_url_path
        else:
            logger.warning("Thumbnail URL path is None, skipping injection")
        
        # 3. Increment Usage (Only if successful)
        if author_id != ADMIN_ID:
            usage_store.increment_used(author_id)
            
        return result

    except httpx.RequestError as exc:
        raise HTTPException(status_code=502, detail=f"Video Inference Service unreachable: {exc}")
    except Exception as exc:
        # If we caught an HTTPException above, re-raise it
        if isinstance(exc, HTTPException):
            raise exc
        raise HTTPException(status_code=500, detail=f"Proxy error: {exc}")

@app.get("/video-status/{task_id}")
async def get_video_status_proxy(task_id: str):
    url = f"{TURBO_DIFFUSION_LOCAL_URL}/status/{task_id}"
    try:
        async with httpx.AsyncClient(timeout=10.0, trust_env=False) as client:
            resp = await client.get(url)
            
        if resp.status_code == 404:
             raise HTTPException(status_code=404, detail="Task not found")
        if resp.status_code != 200:
             raise HTTPException(status_code=resp.status_code, detail=f"Inference service error: {resp.text}")
             
        return resp.json()
    except httpx.RequestError as exc:
        raise HTTPException(status_code=502, detail=f"Video Inference Service unreachable: {exc}")
    except Exception as exc:
        if isinstance(exc, HTTPException):
            raise exc
        raise HTTPException(status_code=500, detail=f"Proxy error: {exc}")

@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:
        if video.thumbnail:
            logger.info(f"Saving video {video.id} with thumbnail: {video.thumbnail}")
        else:
            logger.warning(f"Saving video {video.id} WITHOUT thumbnail")
        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)