app.py

import os
import json
import uuid
import secrets
import zipfile
import io
import httpx
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Any

from fastapi import FastAPI, Depends, HTTPException, status, UploadFile, File, Query, Request
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, FileResponse, RedirectResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from jose import JWTError, jwt
from passlib.context import CryptContext
from pydantic import BaseModel, Field

# --- Configuration ---
# In a real production app, use Hugging Face Space Secrets to set this!
JWT_SECRET_KEY = os.environ.get("JWT_SECRET_KEY", secrets.token_hex(32))
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7  # 7 days

# --- Persistent Data Paths (Hugging Face Spaces use /data for persistent storage) ---
DATA_DIR = "data"
USERS_DB_FILE = os.path.join(DATA_DIR, "users.json")
UPLOAD_DIR = os.path.join(DATA_DIR, "uploads")
STATIC_DIR = "static"

# Create persistent directories if they don't exist
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(STATIC_DIR, exist_ok=True)

# --- Security & Hashing ---
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# --- Pydantic Models (Data Schemas) ---
class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: Optional[str] = None

class WatchHistoryEntry(BaseModel):
    show_id: str
    show_title: str
    season_number: int
    episode_number: int
    watch_timestamp: datetime

class UserBase(BaseModel):
    username: str

class UserCreate(UserBase):
    password: str

class UserInDB(UserBase):
    hashed_password: str
    profile_picture_url: Optional[str] = None
    watch_history: List[Dict[str, Any]] = Field(default_factory=list)

class UserPublic(UserBase):
    profile_picture_url: Optional[str] = None
    watch_history_detailed: Dict[str, Any] = Field(default_factory=dict)
    email: Optional[str] = None

class PasswordChange(BaseModel):
    current_password: str
    new_password: str

# --- Database Helper Functions (using JSON file) ---
def load_users() -> Dict[str, Dict]:
    if not os.path.exists(USERS_DB_FILE):
        return {}
    try:
        with open(USERS_DB_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, FileNotFoundError):
        return {}

def save_users(users_db: Dict[str, Dict]):
    def json_serializer(obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        raise TypeError(f"Type {type(obj)} not serializable")

    with open(USERS_DB_FILE, "w") as f:
        json.dump(users_db, f, indent=4, default=json_serializer)

# --- Password & Token Functions ---
def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password):
    return pwd_context.hash(password)

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, JWT_SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

# --- Dependency to get current user ---
async def get_current_user(token: str = Depends(oauth2_scheme)) -> UserInDB:
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except JWTError:
        raise credentials_exception

    users_db = load_users()
    user_data = users_db.get(token_data.username)
    if user_data is None:
        raise credentials_exception

    return UserInDB(**user_data)


# --- FastAPI App Initialization ---
app = FastAPI(title="Media Auth API")

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

# --- Helper function to structure watch history ---
def structure_watch_history(history_list: List[Dict]) -> Dict:
    structured = {}
    sorted_history = sorted(history_list, key=lambda x: x.get("watch_timestamp", ""), reverse=True)
    
    for item in sorted_history:
        show_id = item.get("show_id")
        show_title = item.get("show_title", "Unknown Show")
        season_num = item.get("season_number")
        episode_num = item.get("episode_number")
        timestamp = item.get("watch_timestamp")

        if not all([show_id, season_num is not None, episode_num is not None, timestamp]):
            continue

        if show_id not in structured:
            structured[show_id] = {
                "show_id": show_id,
                "title": show_title,
                "seasons": {}
            }
        if season_num not in structured[show_id]["seasons"]:
            structured[show_id]["seasons"][season_num] = {
                "season_number": season_num,
                "episodes": {}
            }
        structured[show_id]["seasons"][season_num]["episodes"][episode_num] = timestamp
    return structured

async def get_anime_poster_url(anime_title: str) -> Optional[str]:
    """Fetches the top anime poster URL from Jikan API."""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(f"https://api.jikan.moe/v4/anime?q={anime_title}&limit=1")
            response.raise_for_status()
            data = response.json()
            if data.get("data"):
                return data["data"][0]["images"]["jpg"]["large_image_url"]
    except Exception as e:
        print(f"Error fetching poster for '{anime_title}': {e}")
        return None
    return None

# --- HTML Content for Download UI ---
DOWNLOAD_UI_HTML = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Download Anime Series</title>
    <style>
        @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;700&display=swap');
        body{font-family:'Roboto',sans-serif;background-color:#141414;color:#fff;margin:0;padding:2rem;display:flex;justify-content:center;align-items:center;min-height:100vh}
        .container{background:#1c1c1c;padding:2rem;border-radius:12px;box-shadow:0 10px 30px rgba(0,0,0,0.5);width:100%;max-width:600px}
        h1{color:#E50914;text-align:center;margin-bottom:2rem;font-weight:700}
        .series-list{list-style:none;padding:0;max-height:40vh;overflow-y:auto;border:1px solid #333;border-radius:8px}
        .series-item{display:flex;align-items:center;padding:1rem;border-bottom:1px solid #333;cursor:pointer;transition:background-color 0.2s ease}
        .series-item:last-child{border-bottom:none}
        .series-item:hover{background-color:#2a2a2a}
        .series-item input[type="checkbox"]{margin-right:1rem;width:20px;height:20px;accent-color:#E50914}
        .series-item label{flex-grow:1;font-size:1.1rem}
        .button-container{text-align:center;margin-top:2rem}
        .btn{background-color:#E50914;color:white;border:none;padding:1rem 2rem;font-size:1.2rem;border-radius:8px;cursor:pointer;transition:background-color 0.2s ease;font-weight:700}
        .btn:hover{background-color:#f6121d}
        .btn:disabled{background-color:#555;cursor:not-allowed}
        #loading-overlay{position:fixed;top:0;left:0;width:100%;height:100%;background:rgba(0,0,0,0.85);display:none;flex-direction:column;justify-content:center;align-items:center;z-index:1000}
        .loader{border:8px solid #f3f3f3;border-top:8px solid #E50914;border-radius:50%;width:80px;height:80px;animation:spin 1s linear infinite}
        #loading-text{color:#fff;margin-top:20px;font-size:1.2rem}
        @keyframes spin{0%{transform:rotate(0deg)}100%{transform:rotate(360deg)}}
    </style>
</head>
<body>
    <div id="loading-overlay"><div class="loader"></div><p id="loading-text">Generating your files...</p></div>
    <div class="container">
        <h1>Select Anime to Download</h1>
        <form id="downloadForm">
            <ul id="seriesList" class="series-list"></ul>
            <div class="button-container"><button type="submit" class="btn" id="generateBtn" disabled>Generate Zip</button></div>
        </form>
    </div>

    <script>
        document.addEventListener('DOMContentLoaded', async () => {
            const seriesList = document.getElementById('seriesList');
            const generateBtn = document.getElementById('generateBtn');
            const loadingOverlay = document.getElementById('loading-overlay');
            const apiBaseUrl = ''; // API is at the same origin

            // Get the token from the URL parameters instead of localStorage.
            const urlParams = new URLSearchParams(window.location.search);
            const token = urlParams.get('token');

            if (!token) {
                seriesList.innerHTML = '<li class="series-item" style="justify-content: center; color: #E50914;">Authentication token not found in URL.</li>';
                generateBtn.disabled = true;
                return;
            }

            try {
                // Use the token from the URL for the Bearer authentication header.
                const response = await fetch(`${apiBaseUrl}/users/me`, {
                    headers: { 'Authorization': `Bearer ${token}` }
                });

                if (!response.ok) {
                    throw new Error('Failed to fetch user data. Your session may have expired.');
                }

                const userData = await response.json();
                const history = userData.watch_history_detailed;
                const uniqueSeries = [...new Set(Object.values(history).map(show => show.title))].sort();

                if (uniqueSeries.length === 0) {
                    seriesList.innerHTML = '<li class="series-item" style="justify-content: center;">No watched series found.</li>';
                    return;
                }

                uniqueSeries.forEach(title => {
                    const listItem = document.createElement('li');
                    listItem.className = 'series-item';
                    listItem.innerHTML = `<input type="checkbox" id="${title}" name="series" value="${title}"><label for="${title}">${title}</label>`;
                    seriesList.appendChild(listItem);
                });

                generateBtn.disabled = false;

            } catch (error) {
                seriesList.innerHTML = `<li class="series-item" style="justify-content: center; color: #E50914;">${error.message}</li>`;
            }

            document.getElementById('downloadForm').addEventListener('submit', (e) => {
                e.preventDefault();
                loadingOverlay.style.display = 'flex';

                const selectedSeries = Array.from(document.querySelectorAll('input[name="series"]:checked')).map(cb => cb.value);

                if (selectedSeries.length === 0) {
                    alert('Please select at least one series.');
                    loadingOverlay.style.display = 'none';
                    return;
                }

                const queryString = new URLSearchParams({ series_titles: selectedSeries.join(',') }).toString();
                
                // Construct the download URL, passing the token as a query parameter as required by the /generate-zip endpoint.
                const downloadUrl = `${apiBaseUrl}/generate-zip?${queryString}&token=${token}`;
                
                window.open(downloadUrl, '_blank');
                
                setTimeout(() => { loadingOverlay.style.display = 'none'; }, 3000);
            });
        });
    </script>
</body>
</html>
"""

# --- API Endpoints ---

@app.post("/token", response_model=Token, tags=["Authentication"])
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    users_db = load_users()
    user_data = users_db.get(form_data.username)
    if not user_data or not verify_password(form_data.password, user_data.get("hashed_password")):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user_data["username"]}, expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}


@app.post("/signup", status_code=status.HTTP_201_CREATED, tags=["Authentication"])
async def signup_user(user: UserCreate):
    users_db = load_users()
    if user.username in users_db:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Username already registered",
        )
    hashed_password = get_password_hash(user.password)
    new_user = UserInDB(
        username=user.username,
        hashed_password=hashed_password,
        profile_picture_url=None,
        watch_history=[]
    )
    users_db[user.username] = new_user.dict()
    save_users(users_db)
    return {"message": "User created successfully. Please login."}


@app.get("/users/me", response_model=UserPublic, tags=["User"])
async def read_users_me(current_user: UserInDB = Depends(get_current_user)):
    detailed_history = structure_watch_history(current_user.watch_history)
    
    user_public_data = UserPublic(
        username=current_user.username,
        email=current_user.username,
        profile_picture_url=current_user.profile_picture_url,
        watch_history_detailed=detailed_history
    )
    return user_public_data


@app.post("/users/me/profile-picture", response_model=UserPublic, tags=["User"])
async def upload_profile_picture(
    file: UploadFile = File(...),
    current_user: UserInDB = Depends(get_current_user)
):
    file_extension = os.path.splitext(file.filename)[1].lower()
    if file_extension not in ['.png', '.jpg', '.jpeg', '.gif', '.webp']:
        raise HTTPException(status_code=400, detail="Invalid file type.")
        
    unique_filename = f"{uuid.uuid4()}{file_extension}"
    file_path = os.path.join(UPLOAD_DIR, unique_filename)

    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())

    profile_picture_url = f"/uploads/{unique_filename}"
    users_db = load_users()
    users_db[current_user.username]["profile_picture_url"] = profile_picture_url
    save_users(users_db)

    current_user.profile_picture_url = profile_picture_url
    detailed_history = structure_watch_history(current_user.watch_history)
    return UserPublic(
        username=current_user.username,
        email=current_user.username,
        profile_picture_url=current_user.profile_picture_url,
        watch_history_detailed=detailed_history
    )


@app.get("/users/me/watch-history", status_code=status.HTTP_200_OK, tags=["User"])
async def update_watch_history(
    show_id: str = Query(...),
    show_title: str = Query(...),
    season_number: int = Query(..., ge=0),
    episode_number: int = Query(..., ge=1),
    current_user: UserInDB = Depends(get_current_user)
):
    users_db = load_users()
    user_data = users_db[current_user.username]
    episode_id = f"{show_id}_{season_number}_{episode_number}"

    is_already_watched = any(
        (f"{item.get('show_id')}_{item.get('season_number')}_{item.get('episode_number')}" == episode_id)
        for item in user_data.get("watch_history", [])
    )

    if not is_already_watched:
        new_entry = WatchHistoryEntry(
            show_id=show_id,
            show_title=show_title,
            season_number=season_number,
            episode_number=episode_number,
            watch_timestamp=datetime.utcnow()
        )
        user_data.setdefault("watch_history", []).append(new_entry.dict())
        save_users(users_db)
        return {"message": "Watch history updated."}

    return {"message": "Episode already in watch history."}


@app.post("/users/me/password", status_code=status.HTTP_200_OK, tags=["User"])
async def change_user_password(
    password_data: PasswordChange,
    current_user: UserInDB = Depends(get_current_user)
):
    users_db = load_users()
    user_data = users_db[current_user.username]
    
    if not verify_password(password_data.current_password, user_data["hashed_password"]):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Incorrect current password",
        )
    
    new_hashed_password = get_password_hash(password_data.new_password)
    user_data["hashed_password"] = new_hashed_password
    save_users(users_db)
    
    return {"message": "Password updated successfully"}


@app.get("/download-ui", response_class=HTMLResponse, tags=["Download"])
async def get_download_ui(request: Request):
    """Serves the modern HTML interface for selecting downloads."""
    return HTMLResponse(content=DOWNLOAD_UI_HTML)


@app.get("/generate-zip", tags=["Download"])
async def generate_zip_file(
    series_titles: str = Query(..., description="Comma-separated list of anime titles"),
    token: str = Query(..., description="User's auth token from URL param")
):
    """
    Generates a zip file containing folders for selected anime series,
    each with a poster.png inside. It authenticates using the token from the URL.
    """
    # Create a temporary dependency to validate the token from the query parameter
    async def get_user_from_query_token(token_str: str = token) -> UserInDB:
        # This re-uses the logic from get_current_user but works on a raw token string
        return await get_current_user(token=token_str)

    try:
        # Validate the token by calling our temporary dependency
        await get_user_from_query_token()
    except HTTPException:
        raise HTTPException(status_code=401, detail="Authentication failed from URL token")

    titles = [title.strip() for title in series_titles.split(',') if title.strip()]
    
    zip_buffer = io.BytesIO()

    async with httpx.AsyncClient() as client:
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for title in titles:
                # Sanitize title for folder name
                safe_folder_name = "".join(c for c in title if c.isalnum() or c in " .-_").rstrip()
                folder_path = f"Anime/{safe_folder_name}/"
                
                poster_url = await get_anime_poster_url(title)
                
                if poster_url:
                    try:
                        response = await client.get(poster_url)
                        response.raise_for_status()
                        zipf.writestr(f"{folder_path}poster.png", response.content)
                    except Exception as e:
                        print(f"Failed to download or write poster for '{title}': {e}")
                        # Create an empty file to ensure folder creation
                        zipf.writestr(f"{folder_path}.placeholder", "")
                else:
                    # If no poster is found, still create the folder via a placeholder
                    zipf.writestr(f"{folder_path}.placeholder", "")

    zip_buffer.seek(0)
    
    return StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={"Content-Disposition": "attachment; filename=anime_series_folders.zip"}
    )

# --- Static File Serving & Root Redirect ---
app.mount("/uploads", StaticFiles(directory=UPLOAD_DIR), name="uploads")

@app.get("/", include_in_schema=False)
def root():
    return RedirectResponse(url="/docs")