import os
import numpy as np
from celery_app import app
from config import settings
from models.vision_model import estimate_door_positions
from utils.image_utils import load_equirectangular_image
from utils.geometry import calculate_reverse_yaw, angle_difference
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session


@app.task(bind=True)
def hotspot_prediction(self, house_id: str):
    """Predict hotspot yaw/pitch for detected connections."""
    engine = create_engine(settings.database_url)

    try:
        with Session(engine) as session:
            # Fetch approved connections
            result = session.execute(
                text(
                    "SELECT id, data FROM ai_suggestions "
                    "WHERE house_id = :house_id "
                    "AND suggestion_type = 'connection' "
                    "AND status = 'approved'"
                ),
                {"house_id": house_id},
            )
            suggestions = result.fetchall()

            # Load panorama data
            pano_result = session.execute(
                text(
                    "SELECT id, file_path, capture_order FROM panoramas "
                    "WHERE house_id = :house_id"
                ),
                {"house_id": house_id},
            )
            panoramas = pano_result.fetchall()
            pano_map = {p[0]: {"path": p[1], "order": p[2]} for p in panoramas}

            hotspots = []
            for suggestion_id, data_str in suggestions:
                try:
                    conn = eval(data_str)  # Safe since we wrote it
                except Exception:
                    continue

                source_id = conn.get("from_pano_id")
                target_id = conn.get("to_pano_id")
                suggested_yaw = conn.get("yaw")

                if not source_id or not target_id:
                    continue

                # Use door detection for precise yaw if not already set
                yaw = suggested_yaw
                if yaw is None:
                    source_info = pano_map.get(source_id)
                    if source_info:
                        source_path = os.path.join(
                            settings.upload_dir, source_info["path"]
                        )
                        if os.path.exists(source_path):
                            source_img = np.array(
                                load_equirectangular_image(source_path)
                            )
                            doors = estimate_door_positions(source_img)
                            if doors:
                                yaw = doors[0][0]

                # Determine connection type
                source_info = pano_map.get(source_id)
                target_info = pano_map.get(target_id)
                order_diff = abs(
                    (source_info["order"] if source_info else 0)
                    - (target_info["order"] if target_info else 0)
                )

                if source_info and target_info:
                    if abs(order_diff) == 1:
                        conn_type = "forward"
                    else:
                        conn_type = "door"
                else:
                    conn_type = "door"

                pitch = -20  # Typical human eye level

                hotspot_data = {
                    "from_pano_id": source_id,
                    "to_pano_id": target_id,
                    "yaw": yaw if yaw is not None else 0.0,
                    "pitch": pitch,
                    "type": conn_type,
                    "confidence": conn.get("confidence", 0.5),
                }
                hotspots.append(hotspot_data)

            # Clear old hotspot predictions
            session.execute(
                text(
                    "DELETE FROM ai_suggestions WHERE house_id = :house_id "
                    "AND suggestion_type = 'hotspot'"
                ),
                {"house_id": house_id},
            )

            for hotspot in hotspots:
                session.execute(
                    text(
                        "INSERT INTO ai_suggestions (house_id, suggestion_type, "
                        "data, confidence, status) "
                        "VALUES (:house_id, 'hotspot', :data, :confidence, 'pending')"
                    ),
                    {
                        "house_id": house_id,
                        "data": str(hotspot),
                        "confidence": hotspot["confidence"],
                    },
                )

            session.commit()

    except Exception as e:
        return {"success": False, "hotspots": [], "error": str(e)}

    return {
        "success": True,
        "hotspots": hotspots,
    }