you-music/backend/app/api/auto_download.py

from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from typing import List, Optional
import json
import logging

from app.db.session import get_db
from app.models.models import DownloadJob, Music
from app.schemas.schemas import DownloadJobCreate, DownloadJobResponse
from app.services.search import music_searcher
from app.services.downloader import music_downloader
from app.core.config import settings
import os
from datetime import datetime

router = APIRouter()
logger = logging.getLogger(__name__)


async def process_download_job(job_id: int, db: AsyncSession, direct_url: str = None):
    """Background task to process a download job"""
    try:
        # Get job
        result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
        job = result.scalar_one_or_none()
        if not job:
            return

        # Check if this is a direct URL download (skip search)
        if direct_url:
            logger.info(f"Direct URL download job {job_id}: url='{direct_url}', song_name='{job.song_name}'")
            
            # Skip search, go directly to download
            job.status = "downloading"
            await db.commit()
            
            # Format output filename
            output_filename = job.song_name  # Already formatted as "Artist - Title"
            
            # Download the music
            success, file_path, error = await music_downloader.download_music(
                direct_url,
                output_filename
            )
            
            if success and file_path:
                # Extract metadata
                metadata = await music_downloader.get_music_metadata(file_path)
                
                # Determine source type
                source_type = "youtube" if music_downloader.is_youtube_url(direct_url) else \
                             "bilibili" if music_downloader.is_bilibili_url(direct_url) else "other"
                
                # Get relative path and file format
                relative_path = os.path.relpath(file_path, settings.MUSIC_DIR)
                file_extension = os.path.splitext(file_path)[1][1:]
                
                # Create database record
                db_music = Music(
                    title=metadata.get("title", job.song_name or "Unknown"),
                    artist=metadata.get("artist", "Unknown"),
                    album=metadata.get("album", ""),
                    duration=metadata.get("duration", 0),
                    file_path=relative_path,
                    file_size=os.path.getsize(file_path),
                    file_format=file_extension,
                    file_location=file_path,
                    file_exists=True,
                    source_url=direct_url,
                    source_type=source_type,
                    thumbnail=metadata.get("thumbnail") or job.thumbnail
                )
                
                db.add(db_music)
                await db.commit()
                await db.refresh(db_music)
                
                job.status = "completed"
                job.music_id = db_music.id
                await db.commit()
                
                logger.info(f"Direct download job {job_id} completed: music_id={db_music.id}")
            else:
                job.status = "failed"
                job.error_message = error or "Download failed"
                await db.commit()
                logger.error(f"Direct download job {job_id} failed: {error}")
            
            return

        # Original auto-download flow (with search)
        # Update status to searching
        job.status = "searching"
        await db.commit()

        # Search for the song
        # If artist is included in song_name (format: "Artist - Song"), use it directly
        # Otherwise, just search by song name
        search_results = await music_searcher.search_all(job.song_name, limit=10)

        # Combine results and filter out short clips (< 90 seconds) to avoid samples
        all_results = search_results.get("youtube", []) + search_results.get("bilibili", [])
        
        # Filter out songs shorter than 90 seconds
        filtered_results = [r for r in all_results if r.duration is None or r.duration >= 90]
        
        if not filtered_results:
            # If all results are too short, fall back to original results
            filtered_results = all_results
            if not filtered_results:
                job.status = "failed"
                job.error_message = "No search results found"
                await db.commit()
                return
        
        # Store search results
        job.search_results = json.dumps([r.model_dump() for r in filtered_results])

        # Find priority result (contains "official" or "官方") among filtered results
        priority_index = -1
        for i, result in enumerate(filtered_results):
            title_lower = result.title.lower()
            if "official" in title_lower or "官方" in title_lower:
                priority_index = i
                job.priority = True
                break

        # Select the result (priority if found, otherwise first one)
        selected_index = priority_index if priority_index >= 0 else 0
        selected_result = filtered_results[selected_index]
        job.selected_result = json.dumps(selected_result.model_dump())
        job.selected_result_index = selected_index

        # Check for duplicates by URL
        result_check = await db.execute(
            select(Music).where(Music.source_url == selected_result.url)
        )
        existing_music = result_check.scalar_one_or_none()

        if existing_music:
            # Mark as duplicate and wait for confirmation
            job.is_duplicate = True
            job.duplicate_music_id = existing_music.id
            job.status = "waiting_confirmation"
            await db.commit()
            return

        # Not a duplicate, proceed with download
        job.status = "downloading"
        await db.commit()

        # Format output filename as "artist - title" to avoid downloading wrong songs
        output_filename = f"{selected_result.artist} - {selected_result.title}" if selected_result.artist else selected_result.title
        
        logger.info(f"Auto-download job {job_id}: artist='{selected_result.artist}', title='{selected_result.title}', output_filename='{output_filename}'")
        
        # Download the music
        success, file_path, error = await music_downloader.download_music(
            selected_result.url,
            output_filename
        )

        if success and file_path:
            # Extract metadata
            metadata = await music_downloader.get_music_metadata(file_path)

            # Determine source type
            source_type = "youtube" if music_downloader.is_youtube_url(selected_result.url) else \
                         "bilibili" if music_downloader.is_bilibili_url(selected_result.url) else "other"

            # Get relative path and file format
            relative_path = os.path.relpath(file_path, settings.MUSIC_DIR)
            file_extension = os.path.splitext(file_path)[1][1:]

            # Create database record
            db_music = Music(
                title=metadata.get("title", selected_result.title or "Unknown"),
                artist=metadata.get("artist", selected_result.artist or "Unknown"),
                album=metadata.get("album", ""),
                duration=metadata.get("duration", selected_result.duration or 0),
                file_path=relative_path,
                file_size=os.path.getsize(file_path),
                file_format=file_extension,
                file_location=file_path,
                file_exists=True,
                source_url=selected_result.url,
                source_type=source_type,
                thumbnail=metadata.get("thumbnail") or selected_result.thumbnail
            )

            db.add(db_music)
            await db.commit()
            await db.refresh(db_music)

            # Update job
            job.status = "completed"
            job.music_id = db_music.id
            await db.commit()
        else:
            job.status = "failed"
            job.error_message = error or "Download failed"
            await db.commit()

    except Exception as e:
        # Update job status to failed
        result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
        job = result.scalar_one_or_none()
        if job:
            job.status = "failed"
            job.error_message = str(e)
            await db.commit()


@router.post("/job", response_model=DownloadJobResponse)
async def create_download_job(
    request: DownloadJobCreate,
    background_tasks: BackgroundTasks,
    db: AsyncSession = Depends(get_db)
):
    """Create a new auto-download job"""
    # Format song_name to include artist if provided
    # This ensures better search results
    song_name = request.song_name
    if request.artist and not request.direct_url:
        # Format as "Artist - Song" for better search accuracy (only if not direct URL)
        song_name = f"{request.artist} - {request.song_name}"
    elif request.artist and request.direct_url:
        # For direct URL, still format the output filename
        song_name = f"{request.artist} - {request.song_name}"
    
    if request.direct_url:
        logger.info(f"Creating direct download job: url='{request.direct_url}', filename='{song_name}'")
    else:
        logger.info(f"Creating auto-download job: song_name='{song_name}' (original: '{request.song_name}', artist: '{request.artist}')")
    
    # Create job
    job = DownloadJob(
        song_name=song_name,
        status="pending",
        thumbnail=request.thumbnail
    )

    db.add(job)
    await db.commit()
    await db.refresh(job)

    # Start processing in background
    # Pass direct_url if provided to skip search
    background_tasks.add_task(process_download_job, job.id, db, request.direct_url)

    return job


@router.get("/jobs", response_model=List[DownloadJobResponse])
async def get_download_jobs(
    status: Optional[str] = None,
    skip: int = 0,
    limit: int = 100,
    db: AsyncSession = Depends(get_db)
):
    """Get all download jobs with optional status filter"""
    query = select(DownloadJob).order_by(DownloadJob.created_at.desc())

    if status:
        query = query.where(DownloadJob.status == status)

    result = await db.execute(query.offset(skip).limit(limit))
    jobs = result.scalars().all()
    return jobs


@router.get("/jobs/{job_id}", response_model=DownloadJobResponse)
async def get_download_job(job_id: int, db: AsyncSession = Depends(get_db)):
    """Get a specific download job"""
    result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
    job = result.scalar_one_or_none()

    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    return job


@router.post("/jobs/{job_id}/confirm", response_model=DownloadJobResponse)
async def confirm_download_job(
    job_id: int,
    background_tasks: BackgroundTasks,
    db: AsyncSession = Depends(get_db)
):
    """Confirm a duplicate download job"""
    result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
    job = result.scalar_one_or_none()

    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    if job.status != "waiting_confirmation":
        raise HTTPException(status_code=400, detail="Job is not waiting for confirmation")

    # Mark as confirmed and restart download
    job.confirmed = True
    job.status = "pending"
    job.is_duplicate = False  # Reset duplicate flag to allow download
    await db.commit()

    # Restart processing
    background_tasks.add_task(process_download_job, job.id, db)

    return job


@router.post("/jobs/{job_id}/retry", response_model=DownloadJobResponse)
async def retry_download_job(
    job_id: int,
    background_tasks: BackgroundTasks,
    db: AsyncSession = Depends(get_db)
):
    """Retry a download job (works for failed, cancelled, pending, searching, waiting_confirmation)"""
    result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
    job = result.scalar_one_or_none()

    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    # Can retry any job that's not completed or currently downloading
    if job.status == "completed":
        raise HTTPException(status_code=400, detail="Cannot retry completed job")

    if job.status == "downloading":
        raise HTTPException(status_code=400, detail="Job is currently downloading, cancel it first")

    # Reset job status
    job.status = "pending"
    job.error_message = None
    await db.commit()

    # Restart processing
    background_tasks.add_task(process_download_job, job.id, db)

    return job


@router.post("/jobs/{job_id}/cancel", response_model=DownloadJobResponse)
async def cancel_download_job(
    job_id: int,
    db: AsyncSession = Depends(get_db)
):
    """Cancel a running/pending download job"""
    result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
    job = result.scalar_one_or_none()

    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    # Can only cancel jobs that are not already completed or failed
    if job.status == "completed":
        raise HTTPException(status_code=400, detail="Cannot cancel completed job")

    if job.status == "cancelled":
        raise HTTPException(status_code=400, detail="Job is already cancelled")

    # Mark as cancelled
    job.status = "cancelled"
    job.error_message = "Cancelled by user"
    await db.commit()

    return job


@router.delete("/jobs/{job_id}")
async def delete_download_job(job_id: int, db: AsyncSession = Depends(get_db)):
    """Delete a download job"""
    result = await db.execute(select(DownloadJob).where(DownloadJob.id == job_id))
    job = result.scalar_one_or_none()

    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    await db.delete(job)
    await db.commit()

    return {"message": "Job deleted"}


@router.post("/jobs/clear-completed")
async def clear_completed_jobs(db: AsyncSession = Depends(get_db)):
    """Clear all completed jobs"""
    result = await db.execute(
        select(DownloadJob).where(DownloadJob.status == "completed")
    )
    jobs = result.scalars().all()

    for job in jobs:
        await db.delete(job)

    await db.commit()

    return {"message": f"Cleared {len(jobs)} completed jobs"}