"""Memory service - Time-based queries and aggregations."""

from __future__ import annotations

import logging
import uuid
from datetime import date, datetime, timedelta, timezone
from typing import Optional

from sqlalchemy import and_, cast, desc, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.types import Date

from neuromem.models.memory import Memory
from neuromem.providers.embedding import EmbeddingProvider

logger = logging.getLogger(__name__)


class MemoryService:
    """Service for time-based memory queries and aggregations."""

    def __init__(self, db: AsyncSession, embedding: Optional[EmbeddingProvider] = None):
        self.db = db
        self._embedding = embedding

    async def get_memories_by_time_range(
        self,
        user_id: str,
        start_time: datetime,
        end_time: Optional[datetime] = None,
        memory_type: Optional[str] = None,
        limit: int = 100,
        offset: int = 0,
    ) -> tuple[int, list[Memory]]:
        """Query memories within a time range."""
        if end_time is None:
            end_time = datetime.now(timezone.utc)

        if start_time >= end_time:
            raise ValueError("start_time must be before end_time")

        conditions = [
            Memory.user_id == user_id,
            Memory.created_at >= start_time,
            Memory.created_at < end_time,
        ]

        if memory_type:
            conditions.append(Memory.memory_type == memory_type)

        count_stmt = select(func.count()).select_from(Memory).where(and_(*conditions))
        total = await self.db.scalar(count_stmt) or 0

        stmt = (
            select(Memory)
            .where(and_(*conditions))
            .order_by(desc(Memory.created_at))
            .limit(limit)
            .offset(offset)
        )

        result = await self.db.execute(stmt)
        memories = list(result.scalars().all())

        return total, memories

    async def get_recent_memories(
        self,
        user_id: str,
        days: int = 7,
        memory_types: Optional[list[str]] = None,
        limit: int = 50,
    ) -> list[Memory]:
        """Query memories from the last N days."""
        start_time = datetime.now(timezone.utc) - timedelta(days=days)

        conditions = [
            Memory.user_id == user_id,
            Memory.created_at >= start_time,
        ]

        if memory_types:
            conditions.append(Memory.memory_type.in_(memory_types))

        stmt = (
            select(Memory)
            .where(and_(*conditions))
            .order_by(desc(Memory.created_at))
            .limit(limit)
        )

        result = await self.db.execute(stmt)
        return list(result.scalars().all())

    async def get_daily_memory_stats(
        self,
        user_id: str,
        start_date: date,
        end_date: Optional[date] = None,
    ) -> list[dict]:
        """Get daily memory statistics."""
        if end_date is None:
            end_date = date.today()

        if start_date > end_date:
            raise ValueError("start_date must be before or equal to end_date")

        start_dt = datetime.combine(start_date, datetime.min.time(), tzinfo=timezone.utc)
        end_dt = datetime.combine(end_date, datetime.max.time(), tzinfo=timezone.utc)

        daily_stats = await self.db.execute(
            select(
                cast(Memory.created_at, Date).label("date"),
                func.count(Memory.id).label("count"),
                Memory.memory_type,
            )
            .where(
                and_(
                    Memory.user_id == user_id,
                    Memory.created_at >= start_dt,
                    Memory.created_at <= end_dt,
                )
            )
            .group_by(
                cast(Memory.created_at, Date),
                Memory.memory_type,
            )
            .order_by(cast(Memory.created_at, Date))
        )

        result = {}
        for row in daily_stats:
            date_key = row.date.isoformat()
            if date_key not in result:
                result[date_key] = {"date": row.date, "count": 0, "memory_types": {}}
            result[date_key]["count"] += row.count
            result[date_key]["memory_types"][row.memory_type] = row.count

        return list(result.values())

    async def get_memory_timeline(
        self,
        user_id: str,
        start_date: date,
        end_date: Optional[date] = None,
        granularity: str = "day",
        memory_type: Optional[str] = None,
    ) -> dict:
        """Get memory timeline with aggregation."""
        if end_date is None:
            end_date = date.today()

        if start_date > end_date:
            raise ValueError("start_date must be before or equal to end_date")

        start_dt = datetime.combine(start_date, datetime.min.time(), tzinfo=timezone.utc)
        end_dt = datetime.combine(end_date, datetime.max.time(), tzinfo=timezone.utc)

        trunc_func = func.date_trunc(granularity, Memory.created_at)

        conditions = [
            Memory.user_id == user_id,
            Memory.created_at >= start_dt,
            Memory.created_at <= end_dt,
        ]

        if memory_type:
            conditions.append(Memory.memory_type == memory_type)

        timeline_data = await self.db.execute(
            select(
                trunc_func.label("period"),
                func.count(Memory.id).label("count"),
            )
            .where(and_(*conditions))
            .group_by(trunc_func)
            .order_by(trunc_func)
        )

        data = [
            {
                "period": row.period.isoformat() if row.period else None,
                "count": row.count,
            }
            for row in timeline_data
        ]

        return {
            "user_id": user_id,
            "start_date": start_date,
            "end_date": end_date,
            "granularity": granularity,
            "total_periods": len(data),
            "data": data,
        }

    async def list_all_memories(
        self,
        user_id: str,
        memory_type: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> tuple[int, list[Memory]]:
        """List all memories with optional filtering and pagination.

        Returns:
            Tuple of (total_count, list_of_memories)
        """
        conditions = [Memory.user_id == user_id]

        if memory_type:
            conditions.append(Memory.memory_type == memory_type)

        # Get total count
        count_stmt = select(func.count()).select_from(Memory).where(and_(*conditions))
        total = await self.db.scalar(count_stmt) or 0

        # Get paginated results
        stmt = (
            select(Memory)
            .where(and_(*conditions))
            .order_by(desc(Memory.created_at))
            .limit(limit)
            .offset(offset)
        )

        result = await self.db.execute(stmt)
        memories = list(result.scalars().all())

        return total, memories

    async def get_memory_by_id(
        self,
        memory_id: str | uuid.UUID,
        user_id: str,
    ) -> Optional[Memory]:
        """Get a single memory by ID with ownership check."""
        if isinstance(memory_id, str):
            memory_id = uuid.UUID(memory_id)

        stmt = select(Memory).where(
            and_(Memory.id == memory_id, Memory.user_id == user_id)
        )
        result = await self.db.execute(stmt)
        return result.scalar_one_or_none()

    async def update_memory(
        self,
        memory_id: str | uuid.UUID,
        user_id: str,
        content: Optional[str] = None,
        memory_type: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> Optional[Memory]:
        """Update a memory's content, type, or metadata.

        If content is changed, the embedding vector will be regenerated.

        Args:
            memory_id: UUID of the memory to update
            user_id: User ID (for access control)
            content: New content (triggers embedding regeneration)
            memory_type: New memory type
            metadata: New metadata (replaces existing)

        Returns:
            Updated Memory object or None if not found
        """
        # Fetch existing memory
        memory = await self.get_memory_by_id(memory_id, user_id)
        if not memory:
            return None

        # Update fields
        content_changed = False
        if content is not None and content != memory.content:
            memory.content = content
            content_changed = True

        if memory_type is not None:
            memory.memory_type = memory_type

        if metadata is not None:
            memory.metadata_ = metadata

        # Regenerate embedding if content changed
        if content_changed and self._embedding:
            vector = await self._embedding.embed(content)
            memory.embedding = vector
        elif content_changed and not self._embedding:
            logger.warning(
                "Content changed but no embedding provider available. "
                "Memory vector not updated."
            )

        await self.db.flush()
        return memory

    async def delete_all_memories(
        self,
        user_id: str,
        memory_type: Optional[str] = None,
    ) -> int:
        """Delete all memories for a user, optionally filtered by type.

        Returns:
            Number of deleted memories
        """
        from sqlalchemy import delete

        conditions = [Memory.user_id == user_id]
        if memory_type:
            conditions.append(Memory.memory_type == memory_type)

        stmt = delete(Memory).where(and_(*conditions))
        result = await self.db.execute(stmt)
        await self.db.flush()
        return result.rowcount

    async def delete_memory(
        self,
        memory_id: str | uuid.UUID,
        user_id: str,
    ) -> bool:
        """Delete a memory and its associated data by ID.

        Cascade-deletes related records from memory_history,
        trait_evidence, and memory_sources tables.

        Args:
            memory_id: UUID of the memory to delete
            user_id: User ID (for access control)

        Returns:
            True if deleted, False if not found
        """
        memory = await self.get_memory_by_id(memory_id, user_id)
        if not memory:
            return False

        from sqlalchemy import delete, text

        mid = memory.id
        # Cascade-delete associated records (no FK constraints exist)
        await self.db.execute(text("DELETE FROM memory_history WHERE memory_id = :mid"), {"mid": mid})
        await self.db.execute(text("DELETE FROM trait_evidence WHERE memory_id = :mid"), {"mid": mid})
        await self.db.execute(text("DELETE FROM memory_sources WHERE memory_id = :mid"), {"mid": mid})
        # Also clean trait_evidence where this memory is the trait itself
        await self.db.execute(text("DELETE FROM trait_evidence WHERE trait_id = :mid"), {"mid": mid})

        await self.db.delete(memory)
        await self.db.flush()
        return True

    async def feedback_trait(
        self,
        trait_id: str,
        user_id: str,
        useful: bool,
    ) -> dict:
        """Record user feedback on a trait.

        useful=True: increment reinforcement_count, boost confidence by 0.05 (cap at 1.0)
        useful=False: increment contradiction_count, reduce confidence by 0.05 (floor at 0.0)

        Returns: {trait_id, useful, reinforcement_count, contradiction_count, confidence}
        """
        from sqlalchemy import select

        memory = (await self.db.execute(
            select(Memory).where(Memory.id == trait_id, Memory.user_id == user_id, Memory.memory_type == "trait")
        )).scalar_one_or_none()

        if not memory:
            return None

        current_confidence = 0.5 if memory.trait_confidence is None else memory.trait_confidence
        if useful:
            memory.trait_reinforcement_count = (memory.trait_reinforcement_count or 0) + 1
            memory.trait_confidence = min(1.0, current_confidence + 0.05)
        else:
            memory.trait_contradiction_count = (memory.trait_contradiction_count or 0) + 1
            memory.trait_confidence = max(0.0, current_confidence - 0.05)

        await self.db.flush()

        return {
            "trait_id": str(memory.id),
            "useful": useful,
            "reinforcement_count": memory.trait_reinforcement_count,
            "contradiction_count": memory.trait_contradiction_count,
            "confidence": round(memory.trait_confidence, 4),
        }

    async def find_duplicates(
        self,
        user_id: str,
        threshold: float = 0.92,
        max_pairs: int = 100,
        days_window: int = 30,
    ) -> list[dict]:
        """Find duplicate memory pairs using embedding cosine similarity.

        Only compares recent memories (within days_window) against all memories.
        Same memory_type required. Traits excluded.

        Returns list of dicts: {id1, content1, type1, importance1, id2, content2, type2, importance2, similarity}
        """
        from sqlalchemy import text
        rows = (await self.db.execute(
            text("""
                SELECT
                    m1.id AS id1, m1.content AS content1, m1.memory_type AS type1,
                    m1.importance AS importance1, m1.created_at AS created1,
                    m2.id AS id2, m2.content AS content2, m2.memory_type AS type2,
                    m2.importance AS importance2, m2.created_at AS created2,
                    1 - (m1.embedding <=> m2.embedding) AS similarity
                FROM memories m1
                JOIN memories m2 ON m1.user_id = m2.user_id
                    AND m1.id != m2.id
                    AND m1.memory_type = m2.memory_type
                WHERE m1.user_id = :uid
                    AND m1.created_at >= NOW() - MAKE_INTERVAL(days => :days)
                    AND m1.embedding IS NOT NULL
                    AND m2.embedding IS NOT NULL
                    AND m1.memory_type NOT IN ('trait')
                    AND m2.memory_type NOT IN ('trait')
                    AND 1 - (m1.embedding <=> m2.embedding) > :threshold
                ORDER BY similarity DESC
                LIMIT :max_pairs
            """),
            {"uid": user_id, "threshold": threshold, "max_pairs": max_pairs, "days": days_window},
        )).fetchall()

        return [
            {
                "id1": str(r.id1), "content1": r.content1, "type1": r.type1,
                "importance1": r.importance1 or 0, "created1": r.created1,
                "id2": str(r.id2), "content2": r.content2, "type2": r.type2,
                "importance2": r.importance2 or 0, "created2": r.created2,
                "similarity": round(float(r.similarity), 4),
            }
            for r in rows
        ]

    async def merge_duplicates(
        self,
        user_id: str,
        threshold: float = 0.92,
        max_pairs: int = 100,
        days_window: int = 30,
        dry_run: bool = False,
    ) -> dict:
        """Find and merge duplicate memories. Keeps higher importance (or newer).

        Returns: {duplicates_found, merged_count, kept_ids, deleted_ids}
        """
        pairs = await self.find_duplicates(user_id, threshold, max_pairs, days_window)

        if not pairs:
            return {"duplicates_found": 0, "merged_count": 0, "kept_ids": [], "deleted_ids": []}

        processed = set()
        kept_ids = []
        deleted_ids = []

        for pair in pairs:
            id1, id2 = pair["id1"], pair["id2"]
            if id1 in processed or id2 in processed:
                continue

            if pair["importance1"] > pair["importance2"]:
                keep_id, delete_id = id1, id2
            elif pair["importance2"] > pair["importance1"]:
                keep_id, delete_id = id2, id1
            else:
                # Same importance — keep newer
                if pair["created1"] and pair["created2"] and pair["created1"] >= pair["created2"]:
                    keep_id, delete_id = id1, id2
                else:
                    keep_id, delete_id = id2, id1

            kept_ids.append(keep_id)
            deleted_ids.append(delete_id)
            processed.add(id1)
            processed.add(id2)

        if dry_run:
            return {
                "duplicates_found": len(pairs),
                "merged_count": 0,
                "would_delete": len(deleted_ids),
                "kept_ids": kept_ids,
                "deleted_ids": deleted_ids,
            }

        for did in deleted_ids:
            await self.delete_memory(did, user_id)

        return {
            "duplicates_found": len(pairs),
            "merged_count": len(deleted_ids),
            "kept_ids": kept_ids,
            "deleted_ids": deleted_ids,
        }
