"""Core Flight Processor — Full Processing Pipeline (Stage 10).

Orchestrates: ImageInputPipeline → VO → MetricRefinement → FactorGraph → SSE.
State Machine: NORMAL → LOST → RECOVERY → NORMAL.
"""

from __future__ import annotations

import asyncio
import logging
import time
from enum import Enum
from typing import Optional

import numpy as np

from gps_denied.core.eskf import ESKF
from gps_denied.core.pipeline import ImageInputPipeline
from gps_denied.core.results import ResultManager
from gps_denied.core.sse import SSEEventStreamer
from gps_denied.db.repository import FlightRepository
from gps_denied.schemas import CameraParameters, GPSPoint
from gps_denied.schemas.flight import (
    BatchMetadata,
    BatchResponse,
    BatchUpdateResponse,
    DeleteResponse,
    FlightCreateRequest,
    FlightDetailResponse,
    FlightResponse,
    FlightStatusResponse,
    ObjectGPSResponse,
    UpdateResponse,
    UserFixRequest,
    UserFixResponse,
    Waypoint,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# State Machine
# ---------------------------------------------------------------------------
class TrackingState(str, Enum):
    """Processing state for a flight."""
    NORMAL = "normal"
    LOST = "lost"
    RECOVERY = "recovery"


class FrameResult:
    """Intermediate result of processing a single frame."""

    def __init__(self, frame_id: int):
        self.frame_id = frame_id
        self.gps: Optional[GPSPoint] = None
        self.confidence: float = 0.0
        self.tracking_state: TrackingState = TrackingState.NORMAL
        self.vo_success: bool = False
        self.alignment_success: bool = False


# ---------------------------------------------------------------------------
# FlightProcessor
# ---------------------------------------------------------------------------
class FlightProcessor:
    """Manages business logic, background processing, and frame orchestration."""

    def __init__(
        self,
        repository: FlightRepository,
        streamer: SSEEventStreamer,
        eskf_config=None,
    ) -> None:
        self.repository = repository
        self.streamer = streamer
        self.result_manager = ResultManager(repository, streamer)
        self.pipeline = ImageInputPipeline(storage_dir=".image_storage", max_queue_size=50)
        self._eskf_config = eskf_config  # ESKFConfig or None → default

        # Per-flight processing state
        self._flight_states: dict[str, TrackingState] = {}
        self._prev_images: dict[str, np.ndarray] = {}  # previous frame cache
        self._flight_cameras: dict[str, CameraParameters] = {}  # per-flight camera
        self._altitudes: dict[str, float] = {}          # per-flight altitude (m)
        self._failure_counts: dict[str, int] = {}       # per-flight consecutive failure counter

        # Per-flight ESKF instances (PIPE-01/07)
        self._eskf: dict[str, ESKF] = {}

        # Lazy-initialised component references (set via `attach_components`)
        self._vo = None           # ISequentialVisualOdometry
        self._gpr = None          # IGlobalPlaceRecognition
        self._metric = None       # IMetricRefinement
        self._graph = None        # IFactorGraphOptimizer
        self._recovery = None     # IFailureRecoveryCoordinator
        self._chunk_mgr = None    # IRouteChunkManager
        self._rotation = None     # ImageRotationManager
        self._satellite = None    # SatelliteDataManager (PIPE-02)
        self._coord = None        # CoordinateTransformer (PIPE-02/06)
        self._mavlink = None      # MAVLinkBridge (PIPE-07)

    # ------ Dependency injection for core components ---------
    def attach_components(
        self,
        vo=None,
        gpr=None,
        metric=None,
        graph=None,
        recovery=None,
        chunk_mgr=None,
        rotation=None,
        satellite=None,
        coord=None,
        mavlink=None,
    ):
        """Attach pipeline components after construction (avoids circular deps)."""
        self._vo = vo
        self._gpr = gpr
        self._metric = metric
        self._graph = graph
        self._recovery = recovery
        self._chunk_mgr = chunk_mgr
        self._rotation = rotation
        self._satellite = satellite   # PIPE-02: SatelliteDataManager
        self._coord = coord           # PIPE-02/06: CoordinateTransformer
        self._mavlink = mavlink       # PIPE-07: MAVLinkBridge

    # ------ ESKF lifecycle helpers ----------------------------
    def _init_eskf_for_flight(
        self, flight_id: str, start_gps: GPSPoint, altitude: float
    ) -> None:
        """Create and initialize a per-flight ESKF instance."""
        if flight_id in self._eskf:
            return
        eskf = ESKF(config=self._eskf_config)
        if self._coord:
            try:
                e, n, _ = self._coord.gps_to_enu(flight_id, start_gps)
                eskf.initialize(np.array([e, n, altitude]), time.time())
            except Exception:
                eskf.initialize(np.zeros(3), time.time())
        else:
            eskf.initialize(np.zeros(3), time.time())
        self._eskf[flight_id] = eskf

    def _eskf_to_gps(self, flight_id: str, eskf: ESKF) -> Optional[GPSPoint]:
        """Convert current ESKF ENU position to WGS84 GPS."""
        if not eskf.initialized or self._coord is None:
            return None
        try:
            pos = eskf.position
            return self._coord.enu_to_gps(flight_id, (float(pos[0]), float(pos[1]), float(pos[2])))
        except Exception:
            return None

    # =========================================================
    # process_frame — central orchestration
    # =========================================================
    async def process_frame(
        self,
        flight_id: str,
        frame_id: int,
        image: np.ndarray,
    ) -> FrameResult:
        """
        Process a single UAV frame through the full pipeline.

        State transitions:
            NORMAL  — VO succeeds → ESKF VO update, attempt satellite fix
            LOST    — VO failed → create new chunk, enter RECOVERY
            RECOVERY— try GPR + MetricRefinement → if anchored, merge & return to NORMAL

        PIPE-01: VO result → eskf.update_vo → satellite match → eskf.update_satellite → MAVLink GPS_INPUT
        PIPE-02: SatelliteDataManager + CoordinateTransformer wired for tile selection
        PIPE-04: Consecutive failure counter wired to FailureRecoveryCoordinator
        PIPE-05: ImageRotationManager initialised on first frame
        PIPE-07: ESKF confidence → MAVLink fix_type via bridge.update_state
        """
        result = FrameResult(frame_id)
        state = self._flight_states.get(flight_id, TrackingState.NORMAL)
        eskf = self._eskf.get(flight_id)

        _default_cam = CameraParameters(
            focal_length=4.5, sensor_width=6.17, sensor_height=4.55,
            resolution_width=640, resolution_height=480,
        )

        # ---- PIPE-05: Initialise heading tracking on first frame ----
        if self._rotation and frame_id == 0:
            self._rotation.requires_rotation_sweep(flight_id)  # seeds HeadingHistory

        # ---- 1. Visual Odometry (frame-to-frame) ----
        vo_ok = False
        if self._vo and flight_id in self._prev_images:
            try:
                cam = self._flight_cameras.get(flight_id, _default_cam)
                rel_pose = self._vo.compute_relative_pose(
                    self._prev_images[flight_id], image, cam
                )
                if rel_pose and rel_pose.tracking_good:
                    vo_ok = True
                    result.vo_success = True

                    if self._graph:
                        self._graph.add_relative_factor(
                            flight_id, frame_id - 1, frame_id, rel_pose, np.eye(6)
                        )

                    # PIPE-01: Feed VO relative displacement into ESKF
                    if eskf and eskf.initialized:
                        now = time.time()
                        dt_vo = max(0.01, now - (eskf.last_timestamp or now))
                        eskf.update_vo(rel_pose.translation, dt_vo)
            except Exception as exc:
                logger.warning("VO failed for frame %d: %s", frame_id, exc)

        # Store current image for next frame
        self._prev_images[flight_id] = image

        # ---- PIPE-04: Consecutive failure counter ----
        if not vo_ok and frame_id > 0:
            self._failure_counts[flight_id] = self._failure_counts.get(flight_id, 0) + 1
        else:
            self._failure_counts[flight_id] = 0

        # ---- 2. State Machine transitions ----
        if state == TrackingState.NORMAL:
            if not vo_ok and frame_id > 0:
                state = TrackingState.LOST
                logger.info("Flight %s → LOST at frame %d", flight_id, frame_id)
                if self._recovery:
                    self._recovery.handle_tracking_lost(flight_id, frame_id)

        if state == TrackingState.LOST:
            state = TrackingState.RECOVERY

        if state == TrackingState.RECOVERY:
            recovered = False
            if self._recovery and self._chunk_mgr:
                active_chunk = self._chunk_mgr.get_active_chunk(flight_id)
                if active_chunk:
                    recovered = self._recovery.process_chunk_recovery(
                        flight_id, active_chunk.chunk_id, [image]
                    )
            if recovered:
                state = TrackingState.NORMAL
                result.alignment_success = True
                # PIPE-04: Reset failure count on successful recovery
                self._failure_counts[flight_id] = 0
                logger.info("Flight %s recovered → NORMAL at frame %d", flight_id, frame_id)

        # ---- 3. Satellite position fix (PIPE-01/02) ----
        if state == TrackingState.NORMAL and self._metric:
            sat_tile: Optional[np.ndarray] = None
            tile_bounds = None

            # PIPE-02: Prefer real SatelliteDataManager tiles (ESKF ±3σ selection)
            if self._satellite and eskf and eskf.initialized:
                gps_est = self._eskf_to_gps(flight_id, eskf)
                if gps_est:
                    cov = eskf.covariance
                    sigma_h = float(
                        np.sqrt(np.trace(cov[0:3, 0:3]) / 3.0)
                    ) if cov is not None else 30.0
                    sigma_h = max(sigma_h, 5.0)
                    try:
                        tile_result = await asyncio.get_event_loop().run_in_executor(
                            None,
                            self._satellite.fetch_tiles_for_position,
                            gps_est, sigma_h, 18,
                        )
                        if tile_result:
                            sat_tile, tile_bounds = tile_result
                    except Exception as exc:
                        logger.debug("Satellite tile fetch failed: %s", exc)

            # Fallback: GPR candidate tile (mock image, real bounds)
            if sat_tile is None and self._gpr:
                try:
                    candidates = self._gpr.retrieve_candidate_tiles(image, top_k=1)
                    if candidates:
                        sat_tile = np.zeros((256, 256, 3), dtype=np.uint8)
                        tile_bounds = candidates[0].bounds
                except Exception as exc:
                    logger.debug("GPR tile fallback failed: %s", exc)

            if sat_tile is not None and tile_bounds is not None:
                try:
                    align = self._metric.align_to_satellite(image, sat_tile, tile_bounds)
                    if align and align.matched:
                        result.gps = align.gps_center
                        result.confidence = align.confidence
                        result.alignment_success = True

                        if self._graph:
                            self._graph.add_absolute_factor(
                                flight_id, frame_id,
                                align.gps_center, np.eye(6),
                                is_user_anchor=False,
                            )

                        # PIPE-01: ESKF satellite update — noise from RANSAC confidence
                        if eskf and eskf.initialized and self._coord:
                            try:
                                e, n, _ = self._coord.gps_to_enu(flight_id, align.gps_center)
                                alt = self._altitudes.get(flight_id, 100.0)
                                pos_enu = np.array([e, n, alt])
                                noise_m = 5.0 + 15.0 * (1.0 - float(align.confidence))
                                eskf.update_satellite(pos_enu, noise_m)
                            except Exception as exc:
                                logger.debug("ESKF satellite update failed: %s", exc)
                except Exception as exc:
                    logger.warning("Metric alignment failed at frame %d: %s", frame_id, exc)

        # ---- 4. Graph optimization (incremental) ----
        if self._graph:
            opt_result = self._graph.optimize(flight_id, iterations=5)
            logger.debug(
                "Optimization: converged=%s, error=%.4f",
                opt_result.converged, opt_result.final_error,
            )

        # ---- PIPE-07: Push ESKF state → MAVLink GPS_INPUT ----
        if self._mavlink and eskf and eskf.initialized:
            try:
                eskf_state = eskf.get_state()
                alt = self._altitudes.get(flight_id, 100.0)
                self._mavlink.update_state(eskf_state, altitude_m=alt)
            except Exception as exc:
                logger.debug("MAVLink state push failed: %s", exc)

        # ---- 5. Publish via SSE ----
        result.tracking_state = state
        self._flight_states[flight_id] = state
        await self._publish_frame_result(flight_id, result)
        return result

    async def _publish_frame_result(self, flight_id: str, result: FrameResult):
        """Emit SSE event for processed frame."""
        event_data = {
            "frame_id": result.frame_id,
            "tracking_state": result.tracking_state.value,
            "vo_success": result.vo_success,
            "alignment_success": result.alignment_success,
            "confidence": result.confidence,
        }
        if result.gps:
            event_data["lat"] = result.gps.lat
            event_data["lon"] = result.gps.lon

        await self.streamer.push_event(
            flight_id, event_type="frame_result", data=event_data
        )

    # =========================================================
    # Existing CRUD / REST helpers (unchanged from Stage 3-4)
    # =========================================================
    async def create_flight(self, req: FlightCreateRequest) -> FlightResponse:
        flight = await self.repository.insert_flight(
            name=req.name,
            description=req.description,
            start_lat=req.start_gps.lat,
            start_lon=req.start_gps.lon,
            altitude=req.altitude,
            camera_params=req.camera_params.model_dump(),
        )
        # P0#2: Store camera params for process_frame VO calls
        self._flight_cameras[flight.id] = req.camera_params

        for poly in req.geofences.polygons:
            await self.repository.insert_geofence(
                flight.id,
                nw_lat=poly.north_west.lat,
                nw_lon=poly.north_west.lon,
                se_lat=poly.south_east.lat,
                se_lon=poly.south_east.lon,
            )
        for w in req.rough_waypoints:
            await self.repository.insert_waypoint(flight.id, lat=w.lat, lon=w.lon)

        # Store per-flight altitude for ESKF/pixel projection
        self._altitudes[flight.id] = req.altitude or 100.0

        # PIPE-02: Set ENU origin and initialise ESKF for this flight
        if self._coord:
            self._coord.set_enu_origin(flight.id, req.start_gps)
        self._init_eskf_for_flight(flight.id, req.start_gps, req.altitude or 100.0)

        # Start MAVLink bridge for this flight (origin required for GPS_INPUT)
        if self._mavlink and not self._mavlink._running:
            try:
                asyncio.create_task(self._mavlink.start(req.start_gps))
            except Exception as exc:
                logger.warning("MAVLink bridge start failed: %s", exc)

        return FlightResponse(
            flight_id=flight.id,
            status="prefetching",
            message="Flight created and prefetching started.",
            created_at=flight.created_at,
        )

    async def get_flight(self, flight_id: str) -> FlightDetailResponse | None:
        flight = await self.repository.get_flight(flight_id)
        if not flight:
            return None
        wps = await self.repository.get_waypoints(flight_id)
        state = await self.repository.load_flight_state(flight_id)

        waypoints = [
            Waypoint(
                id=w.id,
                lat=w.lat,
                lon=w.lon,
                altitude=w.altitude,
                confidence=w.confidence,
                timestamp=w.timestamp,
                refined=w.refined,
            )
            for w in wps
        ]

        status = state.status if state else "unknown"
        frames_processed = state.frames_processed if state else 0
        frames_total = state.frames_total if state else 0

        from gps_denied.schemas import Geofences

        return FlightDetailResponse(
            flight_id=flight.id,
            name=flight.name,
            description=flight.description,
            start_gps=GPSPoint(lat=flight.start_lat, lon=flight.start_lon),
            waypoints=waypoints,
            geofences=Geofences(polygons=[]),
            camera_params=flight.camera_params,
            altitude=flight.altitude,
            status=status,
            frames_processed=frames_processed,
            frames_total=frames_total,
            created_at=flight.created_at,
            updated_at=flight.updated_at,
        )

    async def delete_flight(self, flight_id: str) -> DeleteResponse:
        deleted = await self.repository.delete_flight(flight_id)
        # P0#1: Cleanup in-memory state to prevent memory leaks
        self._cleanup_flight(flight_id)
        return DeleteResponse(deleted=deleted, flight_id=flight_id)

    def _cleanup_flight(self, flight_id: str) -> None:
        """Remove all in-memory state for a flight (prevents memory leaks)."""
        self._prev_images.pop(flight_id, None)
        self._flight_states.pop(flight_id, None)
        self._flight_cameras.pop(flight_id, None)
        self._altitudes.pop(flight_id, None)
        self._failure_counts.pop(flight_id, None)
        self._eskf.pop(flight_id, None)
        if self._graph:
            self._graph.delete_flight_graph(flight_id)

    async def update_waypoint(
        self, flight_id: str, waypoint_id: str, waypoint: Waypoint
    ) -> UpdateResponse:
        ok = await self.repository.update_waypoint(
            flight_id,
            waypoint_id,
            lat=waypoint.lat,
            lon=waypoint.lon,
            altitude=waypoint.altitude,
            confidence=waypoint.confidence,
            refined=waypoint.refined,
        )
        return UpdateResponse(updated=ok, waypoint_id=waypoint_id)

    async def batch_update_waypoints(
        self, flight_id: str, waypoints: list[Waypoint]
    ) -> BatchUpdateResponse:
        failed = []
        updated = 0
        for wp in waypoints:
            ok = await self.repository.update_waypoint(
                flight_id,
                wp.id,
                lat=wp.lat,
                lon=wp.lon,
                altitude=wp.altitude,
                confidence=wp.confidence,
                refined=wp.refined,
            )
            if ok:
                updated += 1
            else:
                failed.append(wp.id)
        return BatchUpdateResponse(
            success=(len(failed) == 0), updated_count=updated, failed_ids=failed
        )

    async def queue_images(
        self, flight_id: str, metadata: BatchMetadata, file_count: int
    ) -> BatchResponse:
        state = await self.repository.load_flight_state(flight_id)
        if state:
            total = state.frames_total + file_count
            await self.repository.save_flight_state(
                flight_id, frames_total=total, status="processing"
            )

        next_seq = metadata.end_sequence + 1
        seqs = list(range(metadata.start_sequence, metadata.end_sequence + 1))
        return BatchResponse(
            accepted=True,
            sequences=seqs,
            next_expected=next_seq,
            message=f"Queued {file_count} images.",
        )

    async def handle_user_fix(
        self, flight_id: str, req: UserFixRequest
    ) -> UserFixResponse:
        await self.repository.save_flight_state(
            flight_id, blocked=False, status="processing"
        )

        # Inject operator position into ESKF with high uncertainty (500m)
        eskf = self._eskf.get(flight_id)
        if eskf and eskf.initialized and self._coord:
            try:
                e, n, _ = self._coord.gps_to_enu(flight_id, req.satellite_gps)
                alt = self._altitudes.get(flight_id, 100.0)
                eskf.update_satellite(np.array([e, n, alt]), noise_meters=500.0)
                self._failure_counts[flight_id] = 0
                logger.info("User fix applied for %s: %s", flight_id, req.satellite_gps)
            except Exception as exc:
                logger.warning("User fix ESKF injection failed: %s", exc)

        return UserFixResponse(
            accepted=True, processing_resumed=True, message="Fix applied."
        )

    async def get_flight_status(self, flight_id: str) -> FlightStatusResponse | None:
        state = await self.repository.load_flight_state(flight_id)
        if not state:
            return None
        return FlightStatusResponse(
            status=state.status,
            frames_processed=state.frames_processed,
            frames_total=state.frames_total,
            current_frame=state.current_frame,
            current_heading=None,
            blocked=state.blocked,
            search_grid_size=state.search_grid_size,
            created_at=state.created_at,
            updated_at=state.updated_at,
        )

    async def convert_object_to_gps(
        self, flight_id: str, frame_id: int, pixel: tuple[float, float]
    ) -> ObjectGPSResponse:
        # PIPE-06: Use real CoordinateTransformer + ESKF pose for ray-ground projection
        gps: Optional[GPSPoint] = None
        eskf = self._eskf.get(flight_id)
        if self._coord and eskf and eskf.initialized:
            pos = eskf.position
            quat = eskf.quaternion
            cam = self._flight_cameras.get(flight_id, CameraParameters(
                focal_length=4.5, sensor_width=6.17, sensor_height=4.55,
                resolution_width=640, resolution_height=480,
            ))
            alt = self._altitudes.get(flight_id, 100.0)
            try:
                gps = self._coord.pixel_to_gps(
                    flight_id,
                    pixel,
                    frame_pose={"position": pos},
                    camera_params=cam,
                    altitude=float(alt),
                    quaternion=quat,
                )
            except Exception as exc:
                logger.debug("pixel_to_gps failed: %s", exc)

        # Fallback: return ESKF position projected to ground (no pixel shift)
        if gps is None and eskf:
            gps = self._eskf_to_gps(flight_id, eskf)

        return ObjectGPSResponse(
            gps=gps or GPSPoint(lat=0.0, lon=0.0),
            accuracy_meters=5.0,
            frame_id=frame_id,
            pixel=pixel,
        )

    async def stream_events(self, flight_id: str, client_id: str):
        """Async generator for SSE stream."""
        async for event in self.streamer.stream_generator(flight_id, client_id):
            yield event
