#!/usr/bin/env python3

import sqlite3
import json
import shutil
import hashlib
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, List, Dict, Tuple
import threading
import gzip

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from std_msgs.msg import String
from builtin_interfaces.msg import Time

from saltybot_social_memory.srv import (
    PersonLookup,
    PersonUpdate,
    EncounterRecord,
    DatabaseStats,
)


class RelationshipTier:
    """Relationship tier definitions and promotion logic."""

    TIERS = ["stranger", "regular", "favorite"]
    TIER_THRESHOLDS = {
        "stranger": 0,
        "regular": 5,  # 5+ encounters
        "favorite": 20,  # 20+ encounters or high positive score
    }
    TIER_COOLDOWN = 3600  # 1 hour between promotions

    @staticmethod
    def promote_tier(current_tier: str, encounter_count: int, positive_score: float) -> Tuple[str, bool]:
        """Determine if person should be promoted to higher tier."""
        current_idx = RelationshipTier.TIERS.index(current_tier)

        if current_idx >= len(RelationshipTier.TIERS) - 1:
            return current_tier, False  # Already at max tier

        next_tier = RelationshipTier.TIERS[current_idx + 1]
        threshold = RelationshipTier.TIER_THRESHOLDS[next_tier]

        # Promote if encounter count meets threshold and positive interactions >80%
        if encounter_count >= threshold and positive_score > 0.8:
            return next_tier, True

        return current_tier, False


class SocialMemoryDatabase:
    """SQLite-backed social memory database."""

    def __init__(self, db_path: Path):
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.lock = threading.Lock()
        self._init_schema()

    def _init_schema(self):
        """Initialize database schema if not exists."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS persons (
                    person_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    embedding BLOB,
                    relationship_tier TEXT DEFAULT 'stranger',
                    relationship_score REAL DEFAULT 0.5,
                    interaction_count INTEGER DEFAULT 0,
                    positive_interactions INTEGER DEFAULT 0,
                    greeting_stranger TEXT,
                    greeting_regular TEXT,
                    greeting_favorite TEXT,
                    notes TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_seen TIMESTAMP,
                    last_tier_promotion TIMESTAMP
                )
            """)

            conn.execute("""
                CREATE TABLE IF NOT EXISTS encounters (
                    encounter_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    person_id INTEGER NOT NULL,
                    transcript TEXT,
                    detected_mood TEXT,
                    engagement_score REAL,
                    positive_interaction BOOLEAN,
                    encounter_timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (person_id) REFERENCES persons(person_id) ON DELETE CASCADE
                )
            """)

            conn.execute("""
                CREATE TABLE IF NOT EXISTS backups (
                    backup_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    backup_path TEXT UNIQUE NOT NULL,
                    backup_timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    file_size_mb REAL
                )
            """)

            conn.commit()

    def add_or_update_person(
        self,
        name: str,
        relationship_tier: str = "stranger",
        custom_greetings: Optional[Dict[str, str]] = None,
        notes: str = "",
        embedding: Optional[List[float]] = None,
    ) -> int:
        """Add or update a person profile."""
        with self.lock:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                # Check if person exists
                cursor.execute("SELECT person_id FROM persons WHERE name = ?", (name,))
                result = cursor.fetchone()

                if result:
                    person_id = result[0]
                    # Update existing person
                    cursor.execute("""
                        UPDATE persons SET
                        relationship_tier = ?,
                        greeting_stranger = ?,
                        greeting_regular = ?,
                        greeting_favorite = ?,
                        notes = ?
                        WHERE person_id = ?
                    """, (
                        relationship_tier,
                        custom_greetings.get("stranger", "") if custom_greetings else "",
                        custom_greetings.get("regular", "") if custom_greetings else "",
                        custom_greetings.get("favorite", "") if custom_greetings else "",
                        notes,
                        person_id,
                    ))
                else:
                    # Insert new person
                    embedding_blob = None
                    if embedding:
                        embedding_blob = json.dumps(embedding).encode()

                    cursor.execute("""
                        INSERT INTO persons (
                            name, relationship_tier, greeting_stranger,
                            greeting_regular, greeting_favorite, notes, embedding
                        ) VALUES (?, ?, ?, ?, ?, ?, ?)
                    """, (
                        name,
                        relationship_tier,
                        custom_greetings.get("stranger", "") if custom_greetings else "",
                        custom_greetings.get("regular", "") if custom_greetings else "",
                        custom_greetings.get("favorite", "") if custom_greetings else "",
                        notes,
                        embedding_blob,
                    ))
                    cursor.execute("SELECT last_insert_rowid()")
                    person_id = cursor.fetchone()[0]

                conn.commit()

        return person_id

    def lookup_person(self, query_type: str, query_value: str) -> Optional[Dict]:
        """Lookup a person by ID, name, or recent encounters."""
        with self.lock:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()

                if query_type == "id":
                    cursor.execute("SELECT * FROM persons WHERE person_id = ?", (int(query_value),))
                elif query_type == "name":
                    cursor.execute("SELECT * FROM persons WHERE name LIKE ?", (f"%{query_value}%",))
                elif query_type == "recent":
                    # Get most recent encounters within N days
                    days = int(query_value) if query_value else 7
                    cursor.execute("""
                        SELECT DISTINCT p.* FROM persons p
                        JOIN encounters e ON p.person_id = e.person_id
                        WHERE e.encounter_timestamp > datetime('now', '-' || ? || ' days')
                        ORDER BY p.last_seen DESC
                        LIMIT 5
                    """, (days,))
                else:
                    return None

                row = cursor.fetchone()
                if not row:
                    return None

                # Parse embedding if present
                embedding = None
                if row["embedding"]:
                    embedding = json.loads(row["embedding"])

                return {
                    "person_id": row["person_id"],
                    "name": row["name"],
                    "relationship_tier": row["relationship_tier"],
                    "relationship_score": row["relationship_score"],
                    "interaction_count": row["interaction_count"],
                    "custom_greetings": {
                        "stranger": row["greeting_stranger"] or "",
                        "regular": row["greeting_regular"] or "",
                        "favorite": row["greeting_favorite"] or "",
                    },
                    "notes": row["notes"] or "",
                    "last_seen": row["last_seen"],
                    "embedding": embedding,
                }

    def record_encounter(
        self,
        person_id: int,
        transcript: str,
        mood: str,
        engagement_score: float,
        positive_interaction: bool,
    ) -> Tuple[int, Optional[str]]:
        """Record an encounter and check for tier promotion."""
        with self.lock:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                # Insert encounter
                cursor.execute("""
                    INSERT INTO encounters (
                        person_id, transcript, detected_mood, engagement_score, positive_interaction
                    ) VALUES (?, ?, ?, ?, ?)
                """, (person_id, transcript, mood, engagement_score, positive_interaction))

                cursor.execute("SELECT last_insert_rowid()")
                encounter_id = cursor.fetchone()[0]

                # Update person stats
                cursor.execute("""
                    SELECT interaction_count, positive_interactions, relationship_tier, last_tier_promotion
                    FROM persons WHERE person_id = ?
                """, (person_id,))
                row = cursor.fetchone()
                current_count = row[0]
                positive_count = row[1]
                current_tier = row[2]
                last_promotion = row[3]

                new_interaction_count = current_count + 1
                new_positive_count = positive_count + (1 if positive_interaction else 0)
                positive_score = new_positive_count / max(1, new_interaction_count)

                # Check for tier promotion (with cooldown)
                new_tier = current_tier
                tier_changed = None
                now = datetime.now().isoformat()

                if last_promotion is None or \
                   (datetime.fromisoformat(last_promotion) + timedelta(hours=RelationshipTier.TIER_COOLDOWN)).isoformat() < now:
                    new_tier, promoted = RelationshipTier.promote_tier(
                        current_tier,
                        new_interaction_count,
                        positive_score,
                    )
                    if promoted:
                        tier_changed = new_tier

                # Update person record
                cursor.execute("""
                    UPDATE persons SET
                    interaction_count = ?,
                    positive_interactions = ?,
                    relationship_score = ?,
                    relationship_tier = ?,
                    last_seen = CURRENT_TIMESTAMP,
                    last_tier_promotion = ?
                    WHERE person_id = ?
                """, (
                    new_interaction_count,
                    new_positive_count,
                    positive_score,
                    new_tier,
                    now if tier_changed else None,
                    person_id,
                ))

                conn.commit()

        return encounter_id, tier_changed

    def get_stats(self) -> Dict:
        """Get database statistics."""
        with self.lock:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                cursor.execute("SELECT COUNT(*) FROM persons")
                total_persons = cursor.fetchone()[0]

                cursor.execute("SELECT COUNT(*) FROM encounters")
                total_encounters = cursor.fetchone()[0]

                cursor.execute(
                    "SELECT COUNT(*) FROM persons WHERE relationship_tier = ?",
                    ("favorite",)
                )
                favorites = cursor.fetchone()[0]

                cursor.execute(
                    "SELECT COUNT(*) FROM persons WHERE relationship_tier = ?",
                    ("regular",)
                )
                regulars = cursor.fetchone()[0]

                cursor.execute(
                    "SELECT COUNT(*) FROM persons WHERE relationship_tier = ?",
                    ("stranger",)
                )
                strangers = cursor.fetchone()[0]

                # Get database file size
                db_size_mb = self.db_path.stat().st_size / (1024 * 1024)

                # Get last backup time
                cursor.execute(
                    "SELECT backup_timestamp FROM backups ORDER BY backup_timestamp DESC LIMIT 1"
                )
                last_backup = cursor.fetchone()
                last_backup_time = last_backup[0] if last_backup else None

                return {
                    "total_persons": total_persons,
                    "total_encounters": total_encounters,
                    "favorites": favorites,
                    "regulars": regulars,
                    "strangers": strangers,
                    "database_size_mb": db_size_mb,
                    "last_backup": last_backup_time,
                }

    def backup_database(self, backup_dir: Path) -> Path:
        """Create a gzipped backup of the database."""
        with self.lock:
            backup_dir.mkdir(parents=True, exist_ok=True)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            backup_path = backup_dir / f"social_memory_backup_{timestamp}.db.gz"

            try:
                with open(self.db_path, "rb") as f_in:
                    with gzip.open(backup_path, "wb") as f_out:
                        f_out.write(f_in.read())

                # Record backup in database
                with sqlite3.connect(self.db_path) as conn:
                    cursor = conn.cursor()
                    size_mb = backup_path.stat().st_size / (1024 * 1024)
                    cursor.execute(
                        "INSERT INTO backups (backup_path, file_size_mb) VALUES (?, ?)",
                        (str(backup_path), size_mb),
                    )
                    conn.commit()

                return backup_path
            except Exception as e:
                raise Exception(f"Backup failed: {str(e)}")

    def restore_database(self, backup_path: Path):
        """Restore database from gzipped backup."""
        with self.lock:
            try:
                with gzip.open(backup_path, "rb") as f_in:
                    with open(self.db_path, "wb") as f_out:
                        f_out.write(f_in.read())
                self._init_schema()  # Ensure schema is valid
            except Exception as e:
                raise Exception(f"Restore failed: {str(e)}")


class SocialMemoryNode(Node):
    """ROS2 service server for social memory database."""

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

        # Configuration
        self.declare_parameter("database_path", "/home/seb/saltybot-data/social_memory.db")
        self.declare_parameter("backup_dir", "/home/seb/saltybot-data/backups")
        self.declare_parameter("publish_stats", True)
        self.declare_parameter("stats_publish_interval", 60.0)

        db_path = Path(self.get_parameter("database_path").value)
        backup_dir = Path(self.get_parameter("backup_dir").value)

        # Initialize database
        self.db = SocialMemoryDatabase(db_path)
        self.backup_dir = backup_dir

        # Create service servers
        self.lookup_service = self.create_service(
            PersonLookup,
            "/saltybot/social_memory/lookup",
            self.handle_lookup,
        )

        self.update_service = self.create_service(
            PersonUpdate,
            "/saltybot/social_memory/update",
            self.handle_update,
        )

        self.encounter_service = self.create_service(
            EncounterRecord,
            "/saltybot/social_memory/encounter",
            self.handle_encounter,
        )

        self.stats_service = self.create_service(
            DatabaseStats,
            "/saltybot/social_memory/stats",
            self.handle_stats,
        )

        # Publisher for stats (if enabled)
        if self.get_parameter("publish_stats").value:
            self.stats_pub = self.create_publisher(
                String,
                "/saltybot/social_memory/stats_update",
                10,
            )
            interval = self.get_parameter("stats_publish_interval").value
            self.stats_timer = self.create_timer(interval, self.publish_stats_callback)

        self.get_logger().info(
            f"Social memory node initialized: {db_path}, "
            f"backup dir: {backup_dir}"
        )

    def handle_lookup(self, request: PersonLookup.Request, response: PersonLookup.Response) -> PersonLookup.Response:
        """Handle person lookup requests."""
        try:
            person = self.db.lookup_person(request.query_type, request.query_value)

            if person:
                response.found = True
                response.person_id = person["person_id"]
                response.name = person["name"]
                response.relationship_tier = person["relationship_tier"]
                response.relationship_score = person["relationship_score"]
                response.interaction_count = person["interaction_count"]
                response.notes = person["notes"]

                # Build custom greetings list
                greetings = person["custom_greetings"]
                response.custom_greetings = [
                    greetings.get("stranger", ""),
                    greetings.get("regular", ""),
                    greetings.get("favorite", ""),
                ]

                if person["last_seen"]:
                    last_seen = datetime.fromisoformat(person["last_seen"])
                    response.last_seen.sec = int(last_seen.timestamp())
            else:
                response.found = False

        except Exception as e:
            self.get_logger().error(f"Lookup error: {str(e)}")
            response.found = False

        return response

    def handle_update(self, request: PersonUpdate.Request, response: PersonUpdate.Response) -> PersonUpdate.Response:
        """Handle person update requests."""
        try:
            custom_greetings = {
                "stranger": request.custom_greetings[0] if len(request.custom_greetings) > 0 else "",
                "regular": request.custom_greetings[1] if len(request.custom_greetings) > 1 else "",
                "favorite": request.custom_greetings[2] if len(request.custom_greetings) > 2 else "",
            }

            person_id = self.db.add_or_update_person(
                name=request.name,
                relationship_tier=request.relationship_tier,
                custom_greetings=custom_greetings,
                notes=request.notes,
                embedding=request.embedding if request.embedding else None,
            )

            response.success = True
            response.message = f"Person '{request.name}' updated successfully"
            response.updated_person_id = person_id

        except Exception as e:
            response.success = False
            response.message = f"Update error: {str(e)}"
            self.get_logger().error(f"Update error: {str(e)}")

        return response

    def handle_encounter(self, request: EncounterRecord.Request, response: EncounterRecord.Response) -> EncounterRecord.Response:
        """Handle encounter recording requests."""
        try:
            encounter_id, tier_changed = self.db.record_encounter(
                person_id=request.person_id,
                transcript=request.transcript,
                mood=request.mood,
                engagement_score=request.engagement_score,
                positive_interaction=request.positive_interaction,
            )

            response.success = True
            response.encounter_id = encounter_id
            response.tier_changed_to = tier_changed if tier_changed else ""
            response.message = f"Encounter recorded (ID: {encounter_id})"

            if tier_changed:
                self.get_logger().info(
                    f"Person {request.person_id} promoted to {tier_changed}"
                )

        except Exception as e:
            response.success = False
            response.message = f"Encounter error: {str(e)}"
            self.get_logger().error(f"Encounter error: {str(e)}")

        return response

    def handle_stats(self, request: DatabaseStats.Request, response: DatabaseStats.Response) -> DatabaseStats.Response:
        """Handle statistics requests."""
        try:
            stats = self.db.get_stats()

            response.total_persons = stats["total_persons"]
            response.total_encounters = stats["total_encounters"]
            response.favorites_count = stats["favorites"]
            response.regulars_count = stats["regulars"]
            response.strangers_count = stats["strangers"]
            response.database_size_mb = stats["database_size_mb"]
            response.database_path = str(self.db.db_path)

            if stats["last_backup"]:
                last_backup = datetime.fromisoformat(stats["last_backup"])
                response.last_backup.sec = int(last_backup.timestamp())

        except Exception as e:
            self.get_logger().error(f"Stats error: {str(e)}")

        return response

    def publish_stats_callback(self):
        """Periodically publish database statistics."""
        try:
            stats = self.db.get_stats()
            msg = String()
            msg.data = json.dumps(stats, default=str)
            self.stats_pub.publish(msg)
        except Exception as e:
            self.get_logger().error(f"Stats publish error: {str(e)}")


def main(args=None):
    rclpy.init(args=args)
    node = SocialMemoryNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
