"""Hebbian router — connection weight tracking with SQLite persistence."""

from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from probos.protocols import ConnectionFactory, DatabaseConnection
from probos.types import AgentID

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS hebbian_weights (
    source_id TEXT NOT NULL,
    target_id TEXT NOT NULL,
    rel_type  TEXT NOT NULL DEFAULT 'intent',
    weight    REAL NOT NULL DEFAULT 0.0,
    updated   TEXT NOT NULL,
    PRIMARY KEY (source_id, target_id, rel_type)
);
"""

# Relationship types
REL_INTENT = "intent"  # intent_id → agent_id (Phase 1 default)
REL_AGENT = "agent"  # agent_id → agent_id (Phase 2 verification)
REL_SOCIAL = "social"  # agent_id → agent_id (AD-453 Ward Room interactions)
REL_BUILDER_VARIANT = "builder_variant"  # build_code → native|visiting (AD-353)
REL_STRATEGY = "strategy"  # strategy_id → agent_type (AD-384)

# Key type: (source, target) for backward compat, internally keyed with rel_type
_WeightKey = tuple[AgentID, AgentID]
_FullKey = tuple[AgentID, AgentID, str]


class HebbianRouter:
    """Tracks connection weights between agents.

    Supports two relationship types:
    - intent: intent_id → agent_id (Phase 1 — intent routing)
    - agent: agent_id → agent_id (Phase 2 — verification relationships)

    When Agent A delegates to Agent B and the result is successful,
    the A→B weight increases. Failed interactions cause decay only.
    Weights are persisted to SQLite so learned topology survives restarts.
    """

    def __init__(
        self,
        decay_rate: float = 0.995,
        reward: float = 0.05,
        db_path: str | Path | None = None,
        connection_factory: ConnectionFactory | None = None,
        social_decay_rate: float | None = None,
    ) -> None:
        self.decay_rate = decay_rate
        # AD-571c v1: per-rel_type decay. None -> fall back to decay_rate (v1 default).
        self.social_decay_rate = (
            social_decay_rate if social_decay_rate is not None else decay_rate
        )
        self.reward = reward
        self.db_path = str(db_path) if db_path else None
        # Full key: (source, target, rel_type)
        self._weights: dict[_FullKey, float] = {}
        # Backward-compat view: (source, target) → weight (aggregated)
        self._compat_weights: dict[_WeightKey, float] = {}
        self._db: DatabaseConnection | None = None
        self._connection_factory = connection_factory
        self._tier_registry: Any = None
        # AD-428b v1: skill-weighted routing late-binds. Both default to "off".
        self._skill_service: Any = None
        self._intent_skill_map: dict[str, str] = {}
        if self._connection_factory is None:
            from probos.storage.sqlite_factory import default_factory
            self._connection_factory = default_factory

    def set_tier_registry(self, registry: Any) -> None:
        """Inject agent tier registry for tier-aware reporting (AD-571)."""
        self._tier_registry = registry

    def set_skill_service(self, service: Any) -> None:
        """AD-428b v1: inject AgentSkillService for skill-weighted routing.

        Late-bind. Idempotent. None disables skill weighting (default).
        """
        self._skill_service = service

    def set_intent_skill_map(self, mapping: dict[str, str]) -> None:
        """AD-428b v1: set the intent_id -> skill_id mapping read at routing time.

        Empty mapping (default) means score_with_skill_weight() is a no-op.
        Late-bind. Idempotent. Re-callable.
        """
        # Defensive copy — caller may continue to mutate the original config dict.
        self._intent_skill_map = dict(mapping or {})

    async def score_with_skill_weight(
        self,
        intent_id: "AgentID",
        agent_id: "AgentID",
        base_weight: float,
    ) -> float:
        """AD-428b v1: skill-weighted routing score multiplier.

        Returns base_weight unchanged when:
          - skill_service is not attached
          - intent_skill_map is empty
          - intent_id is not in the map
          - agent has no profile
          - agent has no record for the mapped skill (or it's suspended)

        When all conditions hold, returns:
          base_weight * (1.0 + 0.10 * (proficiency.value - 1))
        capped at 2.0x. FOLLOW (1) is neutral; SHAPE (7) doubles.

        Tier-2 log-and-degrade — any error returns base_weight unchanged.
        """
        if self._skill_service is None or not self._intent_skill_map:
            return base_weight
        skill_id = self._intent_skill_map.get(intent_id)
        if skill_id is None:
            return base_weight
        try:
            profile = await self._skill_service.get_profile(agent_id)
        except Exception:
            logger.debug(
                "AD-428b: skill_service.get_profile failed for %s; returning base_weight",
                agent_id,
                exc_info=True,
            )
            return base_weight
        if profile is None:
            return base_weight
        # Reuse the SkillProfile.get_proficiency helper added in Section 1b.
        try:
            level = profile.get_proficiency(skill_id)
        except AttributeError:
            return base_weight
        if level is None:
            return base_weight
        multiplier = 1.0 + 0.10 * (level.value - 1)
        capped = min(multiplier, 2.0)
        return base_weight * capped

    def _decay_rate_for(self, rel_type: str) -> float:
        """AD-571c: route per-rel_type to the right decay rate."""
        if rel_type == REL_SOCIAL:
            return self.social_decay_rate
        return self.decay_rate

    async def start(self) -> None:
        """Initialize — load weights from SQLite if configured."""
        if self.db_path:
            self._db = await self._connection_factory.connect(self.db_path)
            await self._db.execute("PRAGMA journal_mode=WAL")
            await self._db.execute("PRAGMA busy_timeout=5000")
            await self._db.execute("PRAGMA foreign_keys = ON")
            await self._db.execute(_SCHEMA)
            await self._db.commit()
            await self._load_from_db()
            logger.info(
                "HebbianRouter loaded %d weights from %s",
                len(self._weights),
                self.db_path,
            )

    async def stop(self) -> None:
        """Persist weights and close DB."""
        if self._db:
            await self._save_to_db()
            await self._db.close()
            self._db = None

    def record_interaction(
        self,
        source: AgentID,
        target: AgentID,
        success: bool,
        rel_type: str = REL_INTENT,
    ) -> float:
        """Update connection weight after an interaction.

        Returns the new weight.

        AD-571c v1: utility-utility prune for REL_INTENT only -- when both
        endpoints are AgentTier.UTILITY, the edge is collaborative noise
        (tools don't form intent-routing relationships with each other).
        CORE_INFRASTRUCTURE pairs and any crew-touching edge always record.
        """
        if (
            self._tier_registry is not None
            and rel_type == REL_INTENT
            and self._is_utility_pair(source, target)
        ):
            return 0.0
        full_key = (source, target, rel_type)
        compat_key = (source, target)
        current = self._weights.get(full_key, 0.0)
        rate = self._decay_rate_for(rel_type)

        if success:
            new_weight = current * rate + self.reward
        else:
            new_weight = current * rate

        # Clamp to [0.0, 1.0]
        new_weight = max(0.0, min(1.0, new_weight))
        self._weights[full_key] = new_weight
        self._compat_weights[compat_key] = new_weight
        return new_weight

    def _is_utility_pair(self, source: AgentID, target: AgentID) -> bool:
        """AD-571c: True only when BOTH endpoints are exact-AgentTier.UTILITY.

        Only fires when both source AND target are explicitly registered.
        REL_INTENT source is typically an intent-name string (not an agent_id),
        which would default to UTILITY via get_tier() and incorrectly trigger
        the prune. Requiring registration keeps intent-name sources unaffected.
        """
        from probos.substrate.agent_tier import AgentTier
        try:
            registered = self._tier_registry.all_registered()
            if source not in registered or target not in registered:
                return False
            src_tier = registered[source]
            tgt_tier = registered[target]
        except Exception:
            return False
        return src_tier == AgentTier.UTILITY and tgt_tier == AgentTier.UTILITY

    def record_verification(
        self,
        verifier_id: AgentID,
        target_id: AgentID,
        verified: bool,
    ) -> float:
        """Record an agent-to-agent verification result.

        Convenience method for red team verification interactions.
        """
        return self.record_interaction(
            source=verifier_id,
            target=target_id,
            success=verified,
            rel_type=REL_AGENT,
        )

    def get_weight(
        self,
        source: AgentID,
        target: AgentID,
        rel_type: str | None = None,
    ) -> float:
        if rel_type is not None:
            return self._weights.get((source, target, rel_type), 0.0)
        return self._compat_weights.get((source, target), 0.0)

    def get_preferred_targets(
        self,
        source: AgentID,
        candidates: list[AgentID],
        rel_type: str | None = None,
        hint: str | None = None,           # AD-418
    ) -> list[AgentID]:
        """Rank candidates by their connection weight from source (descending)."""
        if rel_type is not None:
            scored = [
                (c, self._weights.get((source, c, rel_type), 0.0))
                for c in candidates
            ]
        else:
            scored = [
                (c, self._compat_weights.get((source, c), 0.0))
                for c in candidates
            ]

        # AD-418: Boost hinted agent so it wins when all weights are zero
        if hint:
            for i, (agent_id, score) in enumerate(scored):
                if hint in agent_id:
                    scored[i] = (agent_id, score + 1.0)

        scored.sort(key=lambda x: x[1], reverse=True)
        return [c for c, _ in scored]

    def get_agent_weights(self, agent_id: AgentID) -> dict[AgentID, float]:
        """Get all agent-to-agent weights where agent_id is the source."""
        result = {}
        for (src, tgt, rel), w in self._weights.items():
            if src == agent_id and rel == REL_AGENT:
                result[tgt] = w
        return result

    def decay_all(self) -> int:
        """Apply decay to all weights. Returns count of pruned zero-weights.

        AD-571c v1: per-rel_type decay. REL_SOCIAL uses social_decay_rate
        (slow); all other rel_types use decay_rate (existing default).
        """
        pruned = 0
        keys_to_remove = []
        for key, weight in self._weights.items():
            rel_type = key[2]
            new_weight = weight * self._decay_rate_for(rel_type)
            if new_weight < 0.001:
                keys_to_remove.append(key)
                pruned += 1
            else:
                self._weights[key] = new_weight
        for key in keys_to_remove:
            del self._weights[key]
        # Rebuild compat view
        self._compat_weights.clear()
        for (src, tgt, _), w in self._weights.items():
            self._compat_weights[(src, tgt)] = w
        return pruned

    def prune_defunct_agents(self, live_agent_ids: set[str]) -> int:
        """Remove weights referencing agent targets not in the live roster.

        Only prunes entries where the target looks like a pool-based agent ID
        (contains underscore + hex suffix). Intent names are left alone.

        Returns count of pruned entries.
        """
        keys_to_remove = []
        for key in self._weights:
            target = key[1]
            # Heuristic: pool-based agent IDs have format like
            # "poolname_agenttype_N_hexhash" — skip bare intent names
            parts = target.split("_")
            if len(parts) >= 3 and target not in live_agent_ids:
                keys_to_remove.append(key)
        for key in keys_to_remove:
            del self._weights[key]
        # Rebuild compat view
        self._compat_weights.clear()
        for (src, tgt, _), w in self._weights.items():
            self._compat_weights[(src, tgt)] = w
        if keys_to_remove:
            logger.info("Pruned %d Hebbian weights for defunct agents", len(keys_to_remove))
        return len(keys_to_remove)

    @property
    def weight_count(self) -> int:
        return len(self._weights)

    def all_weights(self, crew_only: bool = False) -> dict[tuple[AgentID, AgentID], float]:
        """Backward-compatible: return (source, target) → weight."""
        if crew_only and self._tier_registry:
            return {
                (source, target): weight
                for (source, target), weight in self._compat_weights.items()
                if self._tier_registry.is_crew(source) or self._tier_registry.is_crew(target)
            }
        return dict(self._compat_weights)

    def all_weights_typed(self) -> dict[_FullKey, float]:
        """Return (source, target, rel_type) → weight."""
        return dict(self._weights)

    # ------------------------------------------------------------------
    # SQLite persistence
    # ------------------------------------------------------------------

    async def _load_from_db(self) -> None:
        if not self._db:
            return
        async with self._db.execute(
            "SELECT source_id, target_id, rel_type, weight FROM hebbian_weights"
        ) as cursor:
            async for row in cursor:
                full_key = (row[0], row[1], row[2])
                compat_key = (row[0], row[1])
                self._weights[full_key] = row[3]
                self._compat_weights[compat_key] = row[3]

    async def _save_to_db(self) -> None:
        if not self._db:
            return
        now = datetime.now(timezone.utc).isoformat()
        await self._db.execute("BEGIN IMMEDIATE")
        try:
            await self._db.execute("DELETE FROM hebbian_weights")
            for (source, target, rel_type), weight in self._weights.items():
                await self._db.execute(
                    "INSERT INTO hebbian_weights (source_id, target_id, rel_type, weight, updated) "
                    "VALUES (?, ?, ?, ?, ?)",
                    (source, target, rel_type, weight, now),
                )
            await self._db.commit()
        except Exception:
            await self._db.execute("ROLLBACK")
            raise
        logger.debug("Saved %d hebbian weights to disk", len(self._weights))

    async def save(self) -> None:
        """Manually trigger a save to disk."""
        await self._save_to_db()

    # ------------------------------------------------------------------
    # AD-514: Public API for weight access
    # ------------------------------------------------------------------

    def get_all_weights(self) -> dict[_FullKey, float]:
        """Public API: get all Hebbian weights (copy)."""
        return dict(self.all_weights_typed())

    def set_weight(self, key: _FullKey, value: float) -> None:
        """Set a single routing weight."""
        self._weights[key] = value
        logger.debug("Hebbian weight set: %s = %.4f", key, value)

    def remove_weights_for_agent(self, agent_id: AgentID) -> None:
        """Remove all routing weights involving the given agent."""
        before = len(self._weights)
        self._weights = {k: v for k, v in self._weights.items() if agent_id not in k[:2]}
        removed = before - len(self._weights)
        if removed:
            logger.info("Removed %d Hebbian weights for agent %s", removed, agent_id)

    def get_all_compat_weights(self) -> dict[_WeightKey, float]:
        """Public API: get all compatibility weights (copy)."""
        return dict(self.all_weights())

    def set_compat_weight(self, key: _WeightKey, value: float) -> None:
        """Set a single compatibility weight."""
        self._compat_weights[key] = value
        logger.debug("Compat weight set: %s = %.4f", key, value)

    def remove_compat_weights_for_agent(self, agent_id: AgentID) -> None:
        """Remove all compatibility weights involving the given agent."""
        before = len(self._compat_weights)
        self._compat_weights = {k: v for k, v in self._compat_weights.items() if agent_id not in k}
        removed = before - len(self._compat_weights)
        if removed:
            logger.info("Removed %d compat weights for agent %s", removed, agent_id)
