api_server.py

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
import asyncio
import os
import time
import json
from typing import Optional, Dict, Any, List
from enum import Enum
from pydantic import BaseModel
from rich.progress import (
    Progress,
    SpinnerColumn,
    TimeElapsedColumn,
    DownloadColumn,
    TransferSpeedColumn,
    BarColumn,
    TextColumn,
)
from rich.console import Console
from rich.live import Live
from rich.table import Table
import download_channel

# Initialize rich console for pretty logging
console = Console()

app = FastAPI(title="Telegram Channel Downloader API")

# Track active downloads and their status
active_downloads: Dict[str, Dict[str, Any]] = {}

class FileStatus(str, Enum):
    PENDING = "pending"
    DOWNLOADING = "downloading"
    DOWNLOADED = "downloaded"
    FAILED = "failed"

class ChannelFile(BaseModel):
    message_id: int
    filename: str
    status: FileStatus
    size: Optional[int] = None
    download_time: Optional[float] = None
    error: Optional[str] = None
    upload_path: Optional[str] = None

class DownloadState(BaseModel):
    channel: str
    last_scanned_id: Optional[int] = None
    files: List[ChannelFile] = []
    current_download: Optional[int] = None  # message_id of current download
    last_updated: float = time.time()

class DownloadRequest(BaseModel):
    channel: Optional[str] = None
    message_limit: Optional[int] = None

class DownloadStatus(BaseModel):
    channel: str
    status: str
    message_count: int = 0
    downloaded: int = 0
    downloading: Optional[str] = None
    error: Optional[str] = None

def create_hf_dataset(token: str) -> bool:
    """Create the Hugging Face dataset if it doesn't exist."""
    try:
        from huggingface_hub import create_repo, RepoNotFoundError
        try:
            # Try to create the dataset repository
            create_repo(
                repo_id=download_channel.HF_REPO_ID,
                token=token,
                repo_type="dataset",
                exist_ok=True
            )
            console.print(f"[green]Created or verified dataset:[/green] {download_channel.HF_REPO_ID}")
            
            # Create initial state file
            initial_state = DownloadState(channel=download_channel.CHANNEL)
            with open(download_channel.STATE_FILE, "w", encoding="utf-8") as f:
                json.dump(initial_state.dict(), f, indent=2, ensure_ascii=False)
            
            # Upload initial state
            if download_channel.upload_file_to_hf(
                download_channel.STATE_FILE,
                download_channel.STATE_FILE,
                token
            ):
                console.print("[green]Initialized dataset with empty state file[/green]")
                return True
        except Exception as e:
            console.print(f"[red]Failed to create dataset:[/red] {str(e)}")
            return False
    except ImportError:
        console.print("[red]huggingface_hub not properly installed[/red]")
        return False
    return True

def download_state_from_hf(token: str) -> DownloadState:
    """Try to download the state file from the HF dataset. Returns state dict or creates new."""
    if not token:
        return DownloadState(channel=download_channel.CHANNEL)
    try:
        # Try to download existing state
        local_path = download_channel.hf_hub_download(
            repo_id=download_channel.HF_REPO_ID,
            filename=download_channel.STATE_FILE,
            repo_type="dataset",
            token=token
        )
        with open(local_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            return DownloadState(**data)
    except Exception as e:
        console.print(f"[yellow]No existing state found, creating new dataset:[/yellow] {str(e)}")
        if create_hf_dataset(token):
            console.print("[green]Dataset created successfully![/green]")
            return DownloadState(channel=download_channel.CHANNEL)
        else:
            console.print("[red]Failed to create dataset, using local state only[/red]")
            return DownloadState(channel=download_channel.CHANNEL)

async def clean_downloaded_file(file_path: str):
    """Remove local file after successful upload"""
    try:
        os.remove(file_path)
        console.print(f"[blue]Cleaned up:[/blue] {os.path.basename(file_path)}")
    except Exception as e:
        console.print(f"[yellow]Warning:[/yellow] Could not clean up {file_path}: {e}")

async def update_and_upload_state(state: DownloadState, token: str) -> bool:
    """Update state timestamp and upload to dataset"""
    state.last_updated = time.time()
    try:
        # Save state locally first
        with open(download_channel.STATE_FILE, "w", encoding="utf-8") as f:
            json.dump(state.dict(), f, indent=2, ensure_ascii=False)
        # Upload to dataset
        return download_channel.upload_file_to_hf(
            download_channel.STATE_FILE,
            download_channel.STATE_FILE,
            token
        )
    except Exception as e:
        console.print(f"[red]Failed to update state:[/red] {e}")
        return False

async def process_message(message, state: DownloadState, client) -> Optional[str]:
    """Process a single message, return output path if file downloaded or None"""
    if not message.media:
        return None

    # Check if it's a RAR file
    is_rar = False
    filename = ""
    if message.file:
        filename = getattr(message.file, 'name', '') or ''
        if filename:
            is_rar = filename.lower().endswith('.rar')
        else:
            mime_type = getattr(message.file, 'mime_type', '') or ''
            is_rar = 'rar' in mime_type.lower() if mime_type else False

    if not is_rar:
        return None

    # Use message ID and original filename for saved file
    if filename:
        suggested = f"{message.id}_{filename}"
    else:
        suggested = f"{message.id}.rar"

    return os.path.join(download_channel.OUTPUT_DIR, suggested)

async def run_download(channel: Optional[str], message_limit: Optional[int], task_id: str):
    """Background task to run the download with state management"""
    try:
        # Override channel if provided
        if channel:
            download_channel.CHANNEL = channel
        if message_limit is not None:
            download_channel.MESSAGE_LIMIT = message_limit

        # Get or create download state
        state = download_state_from_hf(download_channel.HF_TOKEN)
        
        # Initialize status for API
        status = {
            "channel": state.channel,
            "status": "running",
            "message_count": len(state.files),
            "downloaded": len([f for f in state.files if f.status == FileStatus.DOWNLOADED]),
            "downloading": None,
            "error": None
        }
        active_downloads[task_id] = status

        # Create progress displays
        progress = Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
            BarColumn(bar_width=40),
            "[progress.percentage]{task.percentage:>3.1f}%",
            "•",
            DownloadColumn(),
            "•",
            TransferSpeedColumn(),
            "•",
            TimeElapsedColumn(),
        )
        
        overall_progress = Progress(
            TextColumn("[bold yellow]{task.description}", justify="right"),
            BarColumn(bar_width=40),
            "[progress.percentage]{task.percentage:>3.1f}%",
            "•",
            TextColumn("[bold green]{task.fields[stats]}")
        )

        # Initialize client
        client = download_channel.TelegramClient(
            download_channel.SESSION_FILE,
            download_channel.API_ID,
            download_channel.API_HASH
        )

        async with client:
            try:
                entity = await client.get_entity(download_channel.CHANNEL)
            except Exception as e:
                console.print(f"[red]Failed to resolve channel:[/red] {e}")
                return 1

            console.print(f"[green]Starting download from:[/green] {entity.title if hasattr(entity, 'title') else download_channel.CHANNEL}")

            # First, scan for new messages and update state
            scan_count = 0
            last_message_id = state.last_scanned_id
            
            try:
                async for message in client.iter_messages(entity, limit=download_channel.MESSAGE_LIMIT or None):
                    scan_count += 1
                    
                    # Update last scanned ID
                    if last_message_id is None or message.id > last_message_id:
                        last_message_id = message.id
                    
                    # Skip if we already know about this message
                    if any(f.message_id == message.id for f in state.files):
                        continue
                        
                    # Check if it's a downloadable file
                    out_path = await process_message(message, state, client)
                    if out_path:
                        # Add to state as pending
                        file_info = ChannelFile(
                            message_id=message.id,
                            filename=os.path.basename(out_path),
                            status=FileStatus.PENDING,
                            size=getattr(message.media, 'size', 0) or 0
                        )
                        state.files.append(file_info)
                
                # Update state with scan results
                state.last_scanned_id = last_message_id
                if download_channel.HF_TOKEN:
                    await update_and_upload_state(state, download_channel.HF_TOKEN)
                    
                console.print(f"[green]Channel scan complete:[/green] Found {scan_count} messages")
                    
            except Exception as e:
                console.print(f"[red]Error during channel scan:[/red] {e}")
                
            # Now process pending downloads
            pending_files = [f for f in state.files if f.status == FileStatus.PENDING]
            total_pending = len(pending_files)

            if total_pending == 0:
                console.print("[green]No new files to download![/green]")
                return 0

            console.print(f"[green]Starting downloads:[/green] {total_pending} files pending")

            # Process pending files
            with Live(progress) as live_progress, Live(overall_progress) as live_overall:
                overall_task = overall_progress.add_task(
                    f"Channel: {download_channel.CHANNEL}",
                    total=total_pending,
                    stats=f"Pending: {total_pending}"
                )

                for file_info in pending_files:
                    try:
                        # Mark as downloading in state
                        file_info.status = FileStatus.DOWNLOADING
                        state.current_download = file_info.message_id
                        if download_channel.HF_TOKEN:
                            await update_and_upload_state(state, download_channel.HF_TOKEN)

                        # Update status
                        status["downloading"] = file_info.filename

                        # Get message and prepare download
                        message = await client.get_messages(entity, ids=file_info.message_id)
                        if not message or not message.media:
                            file_info.status = FileStatus.FAILED
                            file_info.error = "Message not found or no media"
                            continue

                        out_path = os.path.join(download_channel.OUTPUT_DIR, file_info.filename)
                        file_task = progress.add_task(
                            "download",
                            total=file_info.size or 100,
                            filename=file_info.filename
                        )

                        # Download with progress
                        start_time = time.time()
                        try:
                            async def progress_callback(current, total):
                                progress.update(file_task, completed=current)
                                overall_stats = f"Downloaded: {len([f for f in state.files if f.status == FileStatus.DOWNLOADED])}"
                                overall_progress.update(overall_task, completed=current/total*100, stats=overall_stats)

                            await client.download_media(
                                message,
                                file=out_path,
                                progress_callback=progress_callback
                            )

                            # Upload to HF
                            if download_channel.HF_TOKEN:
                                console.print(f"[yellow]Uploading to HF:[/yellow] {file_info.filename}")
                                path_in_repo = f"files/{file_info.filename}"
                                ok = download_channel.upload_file_to_hf(
                                    out_path,
                                    path_in_repo,
                                    download_channel.HF_TOKEN
                                )
                                if ok:
                                    console.print(f"[green]Uploaded:[/green] {file_info.filename}")
                                    # Clean up local file
                                    await clean_downloaded_file(out_path)
                                    file_info.upload_path = path_in_repo
                                else:
                                    console.print(f"[red]Upload failed:[/red] {file_info.filename}")
                                    file_info.error = "Upload to dataset failed"
                                    file_info.status = FileStatus.FAILED
                                    continue

                            # Mark as completed in state
                            file_info.status = FileStatus.DOWNLOADED
                            file_info.download_time = time.time() - start_time
                            
                            # Update state
                            if download_channel.HF_TOKEN:
                                await update_and_upload_state(state, download_channel.HF_TOKEN)

                            # Update status
                            status["downloaded"] += 1
                            await asyncio.sleep(0.2)  # Be polite

                        except download_channel.errors.FloodWaitError as fw:
                            wait = int(fw.seconds) if fw.seconds else 60
                            console.print(f"[yellow]FloodWait:[/yellow] Sleeping {wait}s")
                            await asyncio.sleep(wait + 1)
                            # Retry this file
                            continue

                        except Exception as e:
                            console.print(f"[red]Error:[/red] {str(e)}")
                            file_info.status = FileStatus.FAILED
                            file_info.error = str(e)
                            if download_channel.HF_TOKEN:
                                await update_and_upload_state(state, download_channel.HF_TOKEN)

                    except Exception as e:
                        console.print(f"[red]Fatal error processing {file_info.filename}:[/red] {str(e)}")
                        continue

            # Clear current download
            state.current_download = None
            if download_channel.HF_TOKEN:
                await update_and_upload_state(state, download_channel.HF_TOKEN)

            console.print("[green]Download session completed![/green]")
            status["status"] = "completed"
            status["downloading"] = None

    except Exception as e:
        console.print(f"[red]Fatal error:[/red] {str(e)}")
        if "status" in locals():
            status["status"] = "failed"
            status["error"] = str(e)
    
    return 0

@app.on_event("startup")
async def start_initial_download():
    """Start the download process automatically when the server starts"""
    task_id = "initial_download"
    
    # Verify HF token is set
    if not download_channel.HF_TOKEN:
        console.print("[red]ERROR: HF_TOKEN not set. Please set your Hugging Face token.[/red]")
        return
        
    # Create dataset structure if needed
    console.print("[yellow]Checking Hugging Face dataset...[/yellow]")
    try:
        state = download_state_from_hf(download_channel.HF_TOKEN)
        console.print(f"[green]Using channel:[/green] {state.channel}")
        
        # Create files directory in dataset if it doesn't exist
        os.makedirs(download_channel.OUTPUT_DIR, exist_ok=True)
        
        # Start the download process with default settings
        asyncio.create_task(run_download(
            channel=None,  # Use default from download_channel.py
            message_limit=None,  # Use default
            task_id=task_id
        ))
        console.print(f"[green]Started initial download task:[/green] {task_id}")
        
    except Exception as e:
        console.print(f"[red]Failed to initialize:[/red] {str(e)}")

@app.post("/download", response_model=Dict[str, str])
async def start_download(request: DownloadRequest, background_tasks: BackgroundTasks):
    """Start a new download task"""
    task_id = f"download_{len(active_downloads) + 1}"
    
    background_tasks.add_task(
        run_download,
        channel=request.channel,
        message_limit=request.message_limit,
        task_id=task_id
    )
    
    return {"task_id": task_id}

@app.get("/status/{task_id}", response_model=DownloadStatus)
async def get_status(task_id: str):
    """Get the status of a download task"""
    if task_id not in active_downloads:
        raise HTTPException(status_code=404, detail="Task not found")
    return active_downloads[task_id]

@app.get("/active", response_model=Dict[str, DownloadStatus])
async def list_active():
    """List all active or completed downloads"""
    return active_downloads

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)