"""Procedural memory — directive generation, contradiction detection, confidence decay.

LLM-optimized behavioral directives system that compresses persona +
learned behavioral preferences into cohesive instruction blocks.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from arandu._helpers import cosine_similarity
from arandu.config import MemoryConfig
from arandu.models import MemoryMetaObservation
from arandu.protocols import EmbeddingProvider, LLMProvider

logger = logging.getLogger(__name__)

__all__ = [
    "ContradictionResult",
    "DirectiveBlock",
    "check_directive_contradiction",
    "effective_confidence",
    "generate_optimized_directives",
    "invalidate_directive_cache",
]

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DECAY_RATE_PER_WEEK = 0.95
DECAY_FLOOR = 0.10
META_PROMPT_TIMEOUT_SEC = 5.0
CONTRADICTION_TIMEOUT_SEC = 4.0
DEFAULT_MAX_DIRECTIVES = 10
DEFAULT_MIN_REINFORCEMENTS = 1

# In-memory cache: user_id → (hash, generated_block, generated_at)
_directive_cache: dict[str, tuple[str, str, datetime]] = {}


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------


@dataclass
class DirectiveBlock:
    """Result of directive generation.

    Attributes:
        text: Cohesive behavioral instructions block.
        directive_count: Number of active directives used.
        cache_hit: Whether this was served from cache.
    """

    text: str = ""
    directive_count: int = 0
    cache_hit: bool = False


@dataclass
class ContradictionResult:
    """Result of contradiction check between directives.

    Attributes:
        has_contradiction: Whether a contradiction was found.
        conflicting_directive: Title of the conflicting directive.
        resolution: Explanation of how the contradiction was resolved.
    """

    has_contradiction: bool = False
    conflicting_directive: str | None = None
    resolution: str | None = None


# ---------------------------------------------------------------------------
# Confidence decay (pure function)
# ---------------------------------------------------------------------------


def effective_confidence(
    base_confidence: float,
    created_at: datetime,
    now: datetime | None = None,
) -> float:
    """Apply temporal decay to directive confidence.

    Formula: base_confidence * 0.95^weeks_since_creation.
    ~23% loss after 30 days, ~50% after 90 days.

    Args:
        base_confidence: Original confidence value (0.0-1.0).
        created_at: When the directive was created.
        now: Current timestamp (defaults to UTC now).

    Returns:
        Decayed confidence, floored at 0.10.
    """
    if now is None:
        now = datetime.now(UTC)

    if created_at.tzinfo is None:
        created_at = created_at.replace(tzinfo=UTC)

    days_since = max(0, (now - created_at).total_seconds() / 86400)
    weeks = days_since / 7
    decay = DECAY_RATE_PER_WEEK**weeks
    return float(max(DECAY_FLOOR, base_confidence * decay))


# ---------------------------------------------------------------------------
# Cache helpers
# ---------------------------------------------------------------------------


def _compute_cache_key(prefs: list[Any]) -> str:
    """Hash-based cache key from directive IDs + reinforcement counts."""
    parts = sorted(f"{p.id}:{p.times_reinforced}" for p in prefs)
    raw = "|".join(parts)
    return hashlib.sha256(raw.encode()).hexdigest()[:16]


def invalidate_directive_cache(user_id: str) -> None:
    """Manually invalidate the directive cache for a user."""
    _directive_cache.pop(user_id, None)


def _clear_all_caches() -> None:
    """Clear all cached directive blocks (for testing)."""
    _directive_cache.clear()


# ---------------------------------------------------------------------------
# Meta-prompt
# ---------------------------------------------------------------------------

_META_PROMPT_SYSTEM = """\
You are a prompt engineer specializing in behavioral instructions.

INPUT:
- A persona description (the assistant's base personality)
- A list of learned behavioral directives with confidence scores

TASK:
Generate a SINGLE cohesive block of behavioral instructions that:
1. Integrates persona and directives naturally (no redundancy)
2. Resolves contradictions (persona wins on style, directives win on preferences)
3. Uses imperative, actionable language
4. Prioritizes higher-confidence directives
5. MAXIMUM {budget} tokens — compress without losing meaning

OUTPUT FORMAT (JSON):
{{"directives_block": "The cohesive behavioral instructions text", "directives_used": 3}}

Rules:
- Write in the same language as the persona description. If unclear, use English.
- Do NOT include headers — just the instruction text
- Do NOT number the instructions — use flowing prose with line breaks
- Each instruction should be 1-2 sentences max
"""


async def generate_optimized_directives(
    session: AsyncSession,
    user_id: str,
    llm_provider: LLMProvider,
    config: MemoryConfig,
    *,
    persona_text: str = "",
) -> DirectiveBlock:
    """Generate an LLM-optimized behavioral instructions block.

    Fetches active behavioral_preferences, applies decay, calls LLM
    to produce a cohesive block integrating persona + directives.

    Cache: result cached by hash of directive IDs + reinforcement counts.
    Fail-safe: on any error, returns empty DirectiveBlock.

    Args:
        session: Database session.
        user_id: User identifier.
        llm_provider: Injected LLM provider.
        config: Memory configuration.
        persona_text: Optional persona description.

    Returns:
        DirectiveBlock with generated text.
    """
    # Fetch active behavioral_preferences
    stmt = (
        select(MemoryMetaObservation)
        .where(
            MemoryMetaObservation.user_id == user_id,
            MemoryMetaObservation.observation_type == "behavioral_preference",
            MemoryMetaObservation.is_active.is_(True),
            MemoryMetaObservation.times_reinforced >= DEFAULT_MIN_REINFORCEMENTS,
        )
        .order_by(
            MemoryMetaObservation.times_reinforced.desc(),
            MemoryMetaObservation.last_reinforced_at.desc(),
        )
        .limit(DEFAULT_MAX_DIRECTIVES)
    )
    result = await session.execute(stmt)
    prefs = list(result.scalars().all())

    if not prefs:
        return DirectiveBlock()

    # Check cache
    cache_key = _compute_cache_key(prefs)
    cached = _directive_cache.get(user_id)
    if cached and cached[0] == cache_key:
        # Check TTL
        now = datetime.now(UTC)
        age_minutes = (now - cached[2]).total_seconds() / 60
        if age_minutes < config.directive_cache_ttl_minutes:
            return DirectiveBlock(
                text=cached[1],
                directive_count=len(prefs),
                cache_hit=True,
            )

    # Apply decay and build directive list
    now = datetime.now(UTC)
    directive_lines = []
    for p in prefs:
        eff_conf = effective_confidence(
            p.confidence,
            p.last_reinforced_at or p.created_at,
            now,
        )
        directive_lines.append(f'- "{p.title}" (confidence: {eff_conf:.2f}, reinforced {p.times_reinforced}x)')

    # Build LLM messages
    system_prompt = _META_PROMPT_SYSTEM.format(budget=config.directive_max_tokens)
    user_content = f"PERSONA:\n{persona_text[:500]}\n\nLEARNED DIRECTIVES ({len(prefs)} active):\n" + "\n".join(
        directive_lines
    )

    try:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ]
        _llm_result = await asyncio.wait_for(
            llm_provider.complete(messages),
            timeout=META_PROMPT_TIMEOUT_SEC,
        )
        raw = _llm_result.text
        if not raw:
            return _fallback_block(prefs, now)

        raw = raw.strip()
        if raw.startswith("```"):
            raw = raw.split("\n", 1)[-1] if "\n" in raw else raw[3:]
        if raw.endswith("```"):
            raw = raw.rsplit("```", 1)[0].strip()

        parsed = json.loads(raw)
        block_text = parsed.get("directives_block", "")

        if not block_text:
            return _fallback_block(prefs, now)

        # Cache the result
        generated_at = datetime.now(UTC)
        _directive_cache[user_id] = (cache_key, block_text, generated_at)

        logger.info(
            "generate_optimized_directives: user=%s, directives=%d",
            user_id,
            len(prefs),
        )
        return DirectiveBlock(
            text=block_text,
            directive_count=len(prefs),
            cache_hit=False,
        )

    except Exception as e:
        logger.warning("generate_optimized_directives_failed: user=%s error=%s", user_id, e)
        return _fallback_block(prefs, now)


def _fallback_block(prefs: list[Any], now: datetime) -> DirectiveBlock:
    """Fallback: list of titles as bullet points."""
    lines = [f"- {p.title}" for p in prefs]
    return DirectiveBlock(
        text="\n".join(lines),
        directive_count=len(prefs),
        cache_hit=False,
    )


# ---------------------------------------------------------------------------
# Contradiction detection
# ---------------------------------------------------------------------------

_CONTRADICTION_SYSTEM = """\
You are a behavioral directive analyst. Compare two behavioral directives \
and determine if they CONTRADICT each other.

A contradiction exists when following both directives simultaneously is \
impossible or would produce inconsistent behavior.

OUTPUT (JSON):
{"has_contradiction": true/false, "conflicting_directive": "title", "resolution": "brief explanation"}

Rules:
- Only flag TRUE contradictions (genuinely incompatible behavior)
- When in doubt, say "has_contradiction": false
"""


async def check_directive_contradiction(
    session: AsyncSession,
    user_id: str,
    new_directive: str,
    embedding_provider: EmbeddingProvider,
    llm_provider: LLMProvider,
    *,
    similarity_threshold: float = 0.80,
) -> ContradictionResult:
    """Check a new directive against existing ones for contradictions.

    Uses embedding similarity as pre-filter, then LLM as judge.
    Fail-safe: on error, returns no contradiction.

    Args:
        session: Database session.
        user_id: User identifier.
        new_directive: Text of the new directive to check.
        embedding_provider: Injected embedding provider.
        llm_provider: Injected LLM provider.
        similarity_threshold: Minimum similarity to trigger LLM check.

    Returns:
        ContradictionResult with check outcome.
    """
    if not new_directive or not new_directive.strip():
        return ContradictionResult()

    # Get embedding for new directive
    try:
        new_embedding = await embedding_provider.embed_one(new_directive)
    except Exception as e:
        logger.warning("contradiction_embedding_failed: %s", e)
        return ContradictionResult()

    if not new_embedding:
        return ContradictionResult()

    # Fetch existing active behavioral_preferences
    stmt = select(MemoryMetaObservation).where(
        MemoryMetaObservation.user_id == user_id,
        MemoryMetaObservation.observation_type == "behavioral_preference",
        MemoryMetaObservation.is_active.is_(True),
    )
    result = await session.execute(stmt)
    existing = list(result.scalars().all())

    if not existing:
        return ContradictionResult()

    # Find most similar existing directive
    best_sim = 0.0
    best_existing = None
    for ex in existing:
        ex_vec = getattr(ex, "embedding_vec", None)
        if ex_vec is None:
            continue
        sim = cosine_similarity(new_embedding, list(ex_vec))
        if sim > best_sim:
            best_sim = sim
            best_existing = ex

    if best_sim < similarity_threshold or best_existing is None:
        return ContradictionResult()

    # LLM confirmation
    messages = [
        {"role": "system", "content": _CONTRADICTION_SYSTEM},
        {
            "role": "user",
            "content": (
                f"NEW DIRECTIVE:\n  {new_directive[:200]}\n\n"
                f"EXISTING DIRECTIVE:\n  Title: {best_existing.title}\n"
                f"  Description: {best_existing.text[:200]}\n\n"
                f"Embedding similarity: {best_sim:.3f}"
            ),
        },
    ]
    try:
        _llm_result = await asyncio.wait_for(
            llm_provider.complete(messages),
            timeout=CONTRADICTION_TIMEOUT_SEC,
        )
        raw = _llm_result.text
        if not raw:
            return ContradictionResult()

        raw = raw.strip()
        if raw.startswith("```"):
            raw = raw.split("\n", 1)[-1] if "\n" in raw else raw[3:]
        if raw.endswith("```"):
            raw = raw.rsplit("```", 1)[0].strip()

        parsed = json.loads(raw)
        has_contradiction = bool(parsed.get("has_contradiction", False))
        return ContradictionResult(
            has_contradiction=has_contradiction,
            conflicting_directive=best_existing.title if has_contradiction else None,
            resolution=parsed.get("resolution") if has_contradiction else None,
        )

    except Exception as e:
        logger.warning("contradiction_check_failed: %s", e)
        return ContradictionResult()
