import logging
import asyncio
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession
from typing import List

from app.database import AsyncSessionLocal
from app.services.task_service import task_service
from app.services.db_service import celestial_body_service, position_service
from app.services.horizons import horizons_service

logger = logging.getLogger(__name__)

async def download_positions_task(task_id: int, body_ids: List[str], dates: List[str]):
    """
    Background task worker for downloading NASA positions
    """
    logger.info(f"Task {task_id}: Starting download for {len(body_ids)} bodies and {len(dates)} dates")
    
    async with AsyncSessionLocal() as db:
        try:
            # Mark as running
            await task_service.update_progress(db, task_id, 0, "running")
            
            total_operations = len(body_ids) * len(dates)
            current_op = 0
            success_count = 0
            failed_count = 0
            results = []

            for body_id in body_ids:
                # Check body
                body = await celestial_body_service.get_body_by_id(body_id, db)
                if not body:
                    results.append({"body_id": body_id, "error": "Body not found"})
                    failed_count += len(dates)
                    current_op += len(dates)
                    continue

                body_result = {
                    "body_id": body_id,
                    "body_name": body.name,
                    "dates": []
                }

                for date_str in dates:
                    try:
                        target_date = datetime.strptime(date_str, "%Y-%m-%d")
                        
                        # Check existing
                        existing = await position_service.get_positions(
                            body_id=body_id,
                            start_time=target_date,
                            end_time=target_date.replace(hour=23, minute=59, second=59),
                            session=db
                        )

                        if existing and len(existing) > 0:
                            body_result["dates"].append({"date": date_str, "status": "skipped"})
                            success_count += 1
                        else:
                            # Download
                            positions = horizons_service.get_body_positions(
                                body_id=body_id,
                                start_time=target_date,
                                end_time=target_date,
                                step="1d"
                            )

                            if positions and len(positions) > 0:
                                pos_data = [{
                                    "time": target_date,
                                    "x": positions[0].x,
                                    "y": positions[0].y,
                                    "z": positions[0].z,
                                    "vx": getattr(positions[0], 'vx', None),
                                    "vy": getattr(positions[0], 'vy', None),
                                    "vz": getattr(positions[0], 'vz', None),
                                }]
                                await position_service.save_positions(
                                    body_id=body_id,
                                    positions=pos_data,
                                    source="nasa_horizons",
                                    session=db
                                )
                                body_result["dates"].append({"date": date_str, "status": "success"})
                                success_count += 1
                            else:
                                body_result["dates"].append({"date": date_str, "status": "failed", "error": "No data"})
                                failed_count += 1
                        
                        # Sleep slightly to prevent rate limiting and allow context switching
                        # await asyncio.sleep(0.1) 

                    except Exception as e:
                        logger.error(f"Error processing {body_id} on {date_str}: {e}")
                        body_result["dates"].append({"date": date_str, "status": "error", "error": str(e)})
                        failed_count += 1
                    
                    # Update progress
                    current_op += 1
                    progress = int((current_op / total_operations) * 100)
                    # Only update DB every 5% or so to reduce load, but update Redis frequently
                    # For now, update every item for simplicity
                    await task_service.update_progress(db, task_id, progress)

                results.append(body_result)

            # Complete
            final_result = {
                "total_success": success_count,
                "total_failed": failed_count,
                "details": results
            }
            await task_service.complete_task(db, task_id, final_result)
            logger.info(f"Task {task_id} completed successfully")

        except Exception as e:
            logger.error(f"Task {task_id} failed critically: {e}")
            await task_service.fail_task(db, task_id, str(e))