"""Persistence operations for celestial position data."""

from datetime import datetime
from typing import Any, Dict, List, Optional

from sqlalchemy import and_, delete, func, select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import AsyncSessionLocal
from app.models.db import Position


class PositionService:
    """Store and query time-series position data."""

    @staticmethod
    async def save_positions(
        body_id: str,
        positions: List[Dict[str, Any]],
        source: str = "nasa_horizons",
        session: Optional[AsyncSession] = None,
    ) -> int:
        async def save(active_session: AsyncSession) -> int:
            for position_data in positions:
                statement = insert(Position).values(
                    body_id=body_id,
                    time=position_data["time"],
                    x=position_data["x"],
                    y=position_data["y"],
                    z=position_data["z"],
                    vx=position_data.get("vx"),
                    vy=position_data.get("vy"),
                    vz=position_data.get("vz"),
                    source=source,
                )
                statement = statement.on_conflict_do_update(
                    index_elements=["body_id", "time"],
                    set_={
                        "x": position_data["x"],
                        "y": position_data["y"],
                        "z": position_data["z"],
                        "vx": position_data.get("vx"),
                        "vy": position_data.get("vy"),
                        "vz": position_data.get("vz"),
                        "source": source,
                    },
                )
                await active_session.execute(statement)

            await active_session.commit()
            return len(positions)

        if session is not None:
            return await save(session)
        async with AsyncSessionLocal() as active_session:
            return await save(active_session)

    @staticmethod
    async def get_positions(
        body_id: str,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        session: Optional[AsyncSession] = None,
    ) -> List[Position]:
        async def query(active_session: AsyncSession) -> List[Position]:
            statement = select(Position).where(Position.body_id == body_id)
            if start_time and end_time:
                statement = statement.where(
                    and_(Position.time >= start_time, Position.time <= end_time)
                )
            elif start_time:
                statement = statement.where(Position.time >= start_time)
            elif end_time:
                statement = statement.where(Position.time <= end_time)

            result = await active_session.execute(statement.order_by(Position.time))
            return list(result.scalars().all())

        if session is not None:
            return await query(session)
        async with AsyncSessionLocal() as active_session:
            return await query(active_session)

    @staticmethod
    async def get_positions_in_range(
        body_id: str,
        start_time: datetime,
        end_time: datetime,
        session: Optional[AsyncSession] = None,
    ) -> List[Position]:
        return await PositionService.get_positions(
            body_id, start_time, end_time, session
        )

    @staticmethod
    async def save_position(
        body_id: str,
        time: datetime,
        x: float,
        y: float,
        z: float,
        source: str = "nasa_horizons",
        vx: Optional[float] = None,
        vy: Optional[float] = None,
        vz: Optional[float] = None,
        session: Optional[AsyncSession] = None,
    ) -> Position:
        async def save(active_session: AsyncSession) -> Position:
            result = await active_session.execute(
                select(Position).where(
                    and_(Position.body_id == body_id, Position.time == time)
                )
            )
            position = result.scalar_one_or_none()
            if position is None:
                position = Position(body_id=body_id, time=time, x=x, y=y, z=z)
                active_session.add(position)

            position.x = x
            position.y = y
            position.z = z
            position.vx = vx
            position.vy = vy
            position.vz = vz
            position.source = source
            await active_session.commit()
            await active_session.refresh(position)
            return position

        if session is not None:
            return await save(session)
        async with AsyncSessionLocal() as active_session:
            return await save(active_session)

    @staticmethod
    async def delete_old_positions(
        before_time: datetime,
        session: Optional[AsyncSession] = None,
    ) -> int:
        async def remove(active_session: AsyncSession) -> int:
            result = await active_session.execute(
                delete(Position).where(Position.time < before_time)
            )
            await active_session.commit()
            return result.rowcount

        if session is not None:
            return await remove(session)
        async with AsyncSessionLocal() as active_session:
            return await remove(active_session)

    @staticmethod
    async def get_available_dates(
        body_id: str,
        start_time: datetime,
        end_time: datetime,
        session: Optional[AsyncSession] = None,
    ) -> List[datetime]:
        async def query(active_session: AsyncSession) -> List[datetime]:
            date_expression = func.date(Position.time)
            statement = (
                select(date_expression)
                .where(
                    and_(
                        Position.body_id == body_id,
                        Position.time >= start_time,
                        Position.time <= end_time,
                    )
                )
                .distinct()
                .order_by(date_expression)
            )
            result = await active_session.execute(statement)
            return [row[0] for row in result]

        if session is not None:
            return await query(session)
        async with AsyncSessionLocal() as active_session:
            return await query(active_session)


position_service = PositionService()