#!/usr/bin/env python3
"""First encounter state machine orchestrator.

Detects unknown persons and manages enrollment workflow.
State machine: DETECT → GREET → ASK_NAME → SMALL_TALK → ENROLL → FAREWELL

Subscribes to:
  /saltybot/person_tracker (detected persons with face match confidence)
  /saltybot/stt_result (STT transcriptions)
  /saltybot/person_left (signal when person walks away)

Publishes:
  /social/orchestrator/state (JSON: state, person_id, encounter_data)
  /saltybot/tts_request (Piper TTS triggers)

Stores encounter data as JSON files in /home/seb/encounter-queue/
"""

import json
import time
import threading
from enum import Enum
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Optional, Dict

import rclpy
from rclpy.node import Node
from std_msgs.msg import String, Bool


class EncounterState(Enum):
    """First encounter state machine states."""
    IDLE = "IDLE"
    DETECT = "DETECT"
    GREET = "GREET"
    ASK_NAME = "ASK_NAME"
    SMALL_TALK = "SMALL_TALK"
    ENROLL = "ENROLL"
    FAREWELL = "FAREWELL"


@dataclass
class EncounterData:
    """Encounter metadata and responses."""
    person_id: str
    timestamp: float
    state: str
    name: Optional[str] = None
    context: Optional[str] = None
    greeting_response: Optional[str] = None
    interests: list = None
    enrollment_success: bool = False
    duration_sec: float = 0.0
    notes: str = ""

    def __post_init__(self):
        if self.interests is None:
            self.interests = []


class FirstEncounterOrchestrator(Node):
    """First encounter state machine."""

    # TTS prompts for each state
    TTS_PROMPTS = {
        EncounterState.GREET: "Hi there! I'm SaltyBot. What's your name?",
        EncounterState.ASK_NAME: "Could you please tell me your name?",
        EncounterState.SMALL_TALK: "Nice to meet you! What are you interested in?",
        EncounterState.ENROLL: "I'm saving your face to remember you next time!",
        EncounterState.FAREWELL: "It was great meeting you! Goodbye!",
    }

    def __init__(self):
        super().__init__("first_encounter")

        # Parameters
        self.declare_parameter("encounter_queue_dir", "/home/seb/encounter-queue")
        self.declare_parameter("unknown_face_threshold", 0.3)  # Confidence below threshold = unknown
        self.declare_parameter("small_talk_timeout", 10.0)  # Seconds to wait for STT
        self.declare_parameter("person_away_timeout", 30.0)  # Seconds before person "left"
        self.declare_parameter("auto_enroll", True)

        self.encounter_queue_dir = Path(self.get_parameter("encounter_queue_dir").value)
        self.unknown_threshold = self.get_parameter("unknown_face_threshold").value
        self.small_talk_timeout = self.get_parameter("small_talk_timeout").value
        self.person_away_timeout = self.get_parameter("person_away_timeout").value
        self.auto_enroll = self.get_parameter("auto_enroll").value

        # Create encounter queue directory
        self.encounter_queue_dir.mkdir(parents=True, exist_ok=True)

        # State
        self.current_state = EncounterState.IDLE
        self.current_person_id: Optional[str] = None
        self.encounter_data: Optional[EncounterData] = None
        self.encounter_start_time = 0.0
        self.last_tracker_update = time.time()
        self.state_lock = threading.Lock()
        self.stt_response = None
        self.stt_ready = False

        # Subscriptions
        self.create_subscription(String, "/saltybot/person_tracker", self._on_person_track, 10)
        self.create_subscription(String, "/saltybot/stt_result", self._on_stt_result, 10)
        self.create_subscription(Bool, "/saltybot/person_left", self._on_person_left, 10)

        # Publishers
        self.pub_orchestrator_state = self.create_publisher(String, "/social/orchestrator/state", 10)
        self.pub_tts_request = self.create_publisher(String, "/saltybot/tts_request", 10)

        # Timer for state machine update loop
        self.create_timer(0.5, self._state_machine_update)

        self.get_logger().info(
            f"First encounter orchestrator initialized. Queue: {self.encounter_queue_dir}"
        )

    def _on_person_track(self, msg: String) -> None:
        """Handle person tracker update - detect unknown faces."""
        try:
            data = json.loads(msg.data)
            person_id = data.get("person_id")
            face_confidence = data.get("face_confidence", 1.0)

            # Detect unknown face (low confidence = not in gallery)
            if face_confidence < self.unknown_threshold:
                with self.state_lock:
                    if self.current_state == EncounterState.IDLE:
                        self.current_person_id = person_id
                        self.current_state = EncounterState.DETECT
                        self.encounter_start_time = time.time()
                        self.encounter_data = EncounterData(
                            person_id=person_id,
                            timestamp=self.encounter_start_time,
                            state=EncounterState.DETECT.value
                        )
                        self.get_logger().info(f"Unknown person detected: {person_id}")

            self.last_tracker_update = time.time()
        except json.JSONDecodeError:
            self.get_logger().error(f"Invalid JSON in person tracker: {msg.data}")

    def _on_stt_result(self, msg: String) -> None:
        """Handle STT result."""
        self.stt_response = msg.data
        self.stt_ready = True
        self.get_logger().debug(f"STT result: {msg.data}")

    def _on_person_left(self, msg: Bool) -> None:
        """Handle person walking away - save partial data."""
        if msg.data and self.encounter_data:
            self.get_logger().info(f"Person {self.current_person_id} left. Saving encounter data.")
            self._save_encounter_data("interrupted")
            self._reset_encounter()

    def _send_tts_request(self, text: str) -> None:
        """Send TTS request to Piper."""
        tts_msg = String(data=json.dumps({
            "text": text,
            "voice": "default",
            "speed": 1.0
        }))
        self.pub_tts_request.publish(tts_msg)

    def _publish_state(self) -> None:
        """Publish current orchestrator state."""
        state_data = {
            "state": self.current_state.value,
            "person_id": self.current_person_id,
        }
        if self.encounter_data:
            state_data.update({
                "name": self.encounter_data.name,
                "context": self.encounter_data.context,
            })
        self.pub_orchestrator_state.publish(String(data=json.dumps(state_data)))

    def _wait_for_stt(self, timeout: float) -> Optional[str]:
        """Wait for STT result with timeout."""
        self.stt_ready = False
        self.stt_response = None
        start_time = time.time()

        while time.time() - start_time < timeout:
            if self.stt_ready:
                response = self.stt_response
                self.stt_ready = False
                return response
            time.sleep(0.1)

        return None

    def _state_machine_update(self) -> None:
        """Update state machine."""
        with self.state_lock:
            state = self.current_state

            if state == EncounterState.IDLE:
                pass  # Waiting for detection

            elif state == EncounterState.DETECT:
                # Transition: DETECT → GREET
                self.current_state = EncounterState.GREET
                self._send_tts_request(self.TTS_PROMPTS[EncounterState.GREET])
                self._publish_state()

            elif state == EncounterState.GREET:
                # Wait for STT response (handled asynchronously)
                response = self._wait_for_stt(self.small_talk_timeout)
                if response:
                    self.encounter_data.greeting_response = response
                    self.current_state = EncounterState.ASK_NAME
                    self._send_tts_request(self.TTS_PROMPTS[EncounterState.ASK_NAME])
                    self._publish_state()

            elif state == EncounterState.ASK_NAME:
                # Capture name from STT
                response = self._wait_for_stt(self.small_talk_timeout)
                if response:
                    self.encounter_data.name = response
                    self.current_state = EncounterState.SMALL_TALK
                    self._send_tts_request(self.TTS_PROMPTS[EncounterState.SMALL_TALK])
                    self._publish_state()

            elif state == EncounterState.SMALL_TALK:
                # Capture interests/context
                response = self._wait_for_stt(self.small_talk_timeout)
                if response:
                    self.encounter_data.context = response
                    self.encounter_data.interests = response.split(",")
                    self.current_state = EncounterState.ENROLL
                    self._send_tts_request(self.TTS_PROMPTS[EncounterState.ENROLL])
                    self._publish_state()

            elif state == EncounterState.ENROLL:
                # In real implementation, trigger face enrollment API
                if self.auto_enroll:
                    self.encounter_data.enrollment_success = True
                    self.current_state = EncounterState.FAREWELL
                    self._send_tts_request(self.TTS_PROMPTS[EncounterState.FAREWELL])
                    self._publish_state()

            elif state == EncounterState.FAREWELL:
                # Complete encounter - save data and reset
                self.encounter_data.duration_sec = time.time() - self.encounter_start_time
                self._save_encounter_data("completed")
                self._reset_encounter()

    def _save_encounter_data(self, status: str) -> None:
        """Save encounter data to JSON file."""
        if not self.encounter_data:
            return

        self.encounter_data.notes = status
        self.encounter_data.state = self.current_state.value

        filename = (
            self.encounter_queue_dir /
            f"encounter_{self.encounter_data.person_id}_{int(self.encounter_data.timestamp)}.json"
        )

        try:
            with open(filename, 'w') as f:
                json.dump(asdict(self.encounter_data), f, indent=2)
            self.get_logger().info(f"Encounter data saved: {filename}")
        except Exception as e:
            self.get_logger().error(f"Failed to save encounter data: {e}")

    def _reset_encounter(self) -> None:
        """Reset encounter state."""
        with self.state_lock:
            self.current_state = EncounterState.IDLE
            self.current_person_id = None
            self.encounter_data = None
            self.stt_response = None
            self.stt_ready = False


def main(args=None):
    rclpy.init(args=args)
    node = FirstEncounterOrchestrator()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
