app/routers/migration.py

import os
import shutil
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, status, UploadFile, File
from sqlalchemy.orm import Session
from sqlalchemy.orm.attributes import flag_modified
from pydantic import BaseModel

from app.database import get_db
from app.models.user import User
from app.models.sync import MigrationBackup
from app.routers.auth import get_current_user
from app.utils.hf_storage import (
    upload_to_hf,
    delete_folder_from_hf,
    get_hf_download_url
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/migration", tags=["Device Migration"])

class BackupCreate(BaseModel):
    metadata_json: dict

def cleanup_expired_backups(db: Session, user_id: int):
    """
    Busca y limpia respaldos que hayan pasado el límite de 7 días de validez.
    """
    now = datetime.utcnow()
    expired = db.query(MigrationBackup).filter(
        MigrationBackup.user_id == user_id,
        MigrationBackup.expires_at < now,
        MigrationBackup.status.in_(["pending", "completed"])
    ).all()
    
    for eb in expired:
        try:
            logger.info(f"Limpiando respaldo expirado {eb.id} para usuario {user_id}")
            delete_folder_from_hf(f"migrations/{user_id}")
            eb.status = "expired"
        except Exception as e:
            logger.error(f"Error limpiando respaldo expirado {eb.id}: {e}")
            
    if expired:
        db.commit()

@router.post("/backup")
def create_backup(
    data: BackupCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Inicia un proceso de respaldo. Elimina cualquier respaldo anterior en Hugging Face y BD.
    """
    # 1. Limpieza de expirados
    cleanup_expired_backups(db, current_user.id)
    
    # 2. Buscar si ya existe algún respaldo activo (pendiente o completado)
    active_backups = db.query(MigrationBackup).filter(
        MigrationBackup.user_id == current_user.id,
        MigrationBackup.status.in_(["pending", "completed"])
    ).all()
    
    # Eliminar archivos en Hugging Face y registros en BD de respaldos anteriores
    if active_backups:
        logger.info(f"Eliminando {len(active_backups)} respaldos activos previos para usuario {current_user.id}")
        delete_folder_from_hf(f"migrations/{current_user.id}")
        for ab in active_backups:
            db.delete(ab)
        db.commit()
        
    # 3. Crear el nuevo registro
    expires_at = datetime.utcnow() + timedelta(days=7)
    backup = MigrationBackup(
        user_id=current_user.id,
        expires_at=expires_at,
        status="pending",
        metadata_json=data.metadata_json,
        files_info=[]
    )
    db.add(backup)
    db.commit()
    db.refresh(backup)
    
    return {
        "success": True,
        "backup_id": backup.id,
        "expires_at": backup.expires_at.isoformat()
    }

@router.post("/upload/{backup_id}/{task_id}/{stem_name}")
def upload_backup_file(
    backup_id: str,
    task_id: str,
    stem_name: str,
    file: UploadFile = File(...),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Recibe un stem individual en formato .flac o .json y lo sube a la ruta migrations/ de Hugging Face.
    """
    # Buscar y validar el respaldo
    backup = db.query(MigrationBackup).filter(
        MigrationBackup.id == backup_id,
        MigrationBackup.user_id == current_user.id
    ).first()
    
    if not backup:
        raise HTTPException(status_code=404, detail="Respaldo de migración no encontrado")
        
    if backup.status == "restored":
        raise HTTPException(status_code=400, detail="Este respaldo ya ha sido restaurado")
        
    if backup.expires_at < datetime.utcnow():
        backup.status = "expired"
        db.commit()
        raise HTTPException(status_code=400, detail="El respaldo de migración ha expirado")
        
    # Crear directorio temporal local para recibir el archivo
    temp_dir = os.path.join(os.getcwd(), "uploads", "migration_temp")
    os.makedirs(temp_dir, exist_ok=True)
    orig_ext = ""
    if file.filename:
        orig_ext = file.filename.split('.')[-1].lower()

    if stem_name.endswith("_pitch"):
        ext = "json"
    elif orig_ext in ["m4a", "flac"]:
        ext = orig_ext
    else:
        ext = "flac"
        
    temp_file_path = os.path.join(temp_dir, f"{backup_id}_{task_id}_{stem_name}.{ext}")
    
    try:
        # Guardar en disco local de forma temporal
        with open(temp_file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        # Subir a Hugging Face Dataset
        object_key = f"migrations/{current_user.id}/{task_id}/{stem_name}.{ext}"
        public_url = upload_to_hf(temp_file_path, object_key)
        
        # Registrar el archivo en la BD
        files_list = list(backup.files_info) if backup.files_info else []
        files_list.append({
            "task_id": task_id,
            "stem_name": stem_name,
            "file_key": object_key
        })
        backup.files_info = files_list
        flag_modified(backup, "files_info")
        db.commit()
        
    except Exception as e:
        logger.error(f"Error subiendo archivo de migración {stem_name} para {task_id}: {e}")
        raise HTTPException(status_code=500, detail=f"Error en la subida a Hugging Face: {str(e)}")
    finally:
        if os.path.exists(temp_file_path):
            os.remove(temp_file_path)
            
    return {"success": True, "url": public_url}

@router.post("/finalize/{backup_id}")
def finalize_backup(
    backup_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Marca el respaldo de migración como completado y listo para ser descargado.
    """
    backup = db.query(MigrationBackup).filter(
        MigrationBackup.id == backup_id,
        MigrationBackup.user_id == current_user.id
    ).first()
    
    if not backup:
        raise HTTPException(status_code=404, detail="Respaldo no encontrado")
        
    if backup.expires_at < datetime.utcnow():
        backup.status = "expired"
        db.commit()
        raise HTTPException(status_code=400, detail="El respaldo de migración ha expirado")
        
    backup.status = "completed"
    db.commit()
    
    return {"success": True}

@router.get("/status")
def get_backup_status(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Verifica si existe un respaldo completado y vigente para el usuario actual.
    """
    cleanup_expired_backups(db, current_user.id)
    
    backup = db.query(MigrationBackup).filter(
        MigrationBackup.user_id == current_user.id,
        MigrationBackup.status == "completed"
    ).order_by(MigrationBackup.created_at.desc()).first()
    
    if not backup:
        return {"backup_available": False}
        
    return {
        "backup_available": True,
        "backup_id": backup.id,
        "created_at": backup.created_at.isoformat(),
        "expires_at": backup.expires_at.isoformat(),
        "files_count": len(backup.files_info) if backup.files_info else 0
    }

@router.get("/restore")
def restore_backup(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Retorna el mapa de AsyncStorage y las URLs de descarga de Hugging Face para todos los stems respaldados.
    """
    cleanup_expired_backups(db, current_user.id)
    
    backup = db.query(MigrationBackup).filter(
        MigrationBackup.user_id == current_user.id,
        MigrationBackup.status == "completed"
    ).order_by(MigrationBackup.created_at.desc()).first()
    
    if not backup:
        raise HTTPException(status_code=404, detail="No tienes ningún respaldo activo listo para restaurar")
        
    # Generar URLs de descarga de Hugging Face
    restore_files = []
    files_list = backup.files_info or []
    
    for f in files_list:
        file_key = f["file_key"]
        download_url = get_hf_download_url(file_key)
        restore_files.append({
            "task_id": f["task_id"],
            "stem_name": f["stem_name"],
            "url": download_url,
            "file_key": file_key
        })
        
    return {
        "metadata_json": backup.metadata_json,
        "files": restore_files
    }

@router.post("/complete")
def complete_restore(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Confirma que el nuevo dispositivo ha terminado de descargar los archivos y elimina los datos de Hugging Face.
    """
    backup = db.query(MigrationBackup).filter(
        MigrationBackup.user_id == current_user.id,
        MigrationBackup.status == "completed"
    ).order_by(MigrationBackup.created_at.desc()).first()
    
    if backup:
        # Eliminar archivos físicamente de Hugging Face
        logger.info(f"Limpiando carpeta de HF tras restauración exitosa de {backup.id}")
        delete_folder_from_hf(f"migrations/{current_user.id}")
        backup.status = "restored"
        db.commit()
        
    return {"success": True}

@router.delete("/backup")
def delete_backup(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """
    Elimina permanentemente cualquier respaldo activo de migración del usuario en Hugging Face y la base de datos.
    """
    active_backups = db.query(MigrationBackup).filter(
        MigrationBackup.user_id == current_user.id,
        MigrationBackup.status.in_(["pending", "completed"])
    ).all()
    
    if active_backups:
        logger.info(f"Eliminando {len(active_backups)} respaldos activos por petición de limpieza para usuario {current_user.id}")
        try:
            delete_folder_from_hf(f"migrations/{current_user.id}")
        except Exception as e:
            logger.error(f"Error borrando carpeta de HF para migración del usuario {current_user.id}: {e}")
        for ab in active_backups:
            db.delete(ab)
        db.commit()
        
    return {"success": True}