"""
Constitution Service - Dynamic AI agent context composition.

Composes personalized agent context from multiple layers:
1. Base Constitution (YAML) - LightWave core principles
2. User Identity (Django) - Legal/financial/professional context
3. User Principles (Django) - Personal values and preferences
4. Second Brain (Django) - Weighted knowledge with semantic search

Priority: User overrides > User additions > Base defaults
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from lightwave.ai.constitution.formatters import format_constitution_prompt
from lightwave.schema.pydantic.models.constitution import (
    AgentDefaults,
    ComposedConstitution,
    ConstitutionConfig,
    KnowledgeBook,
    Principle,
    SecondBrainEntryContext,
    UserIdentityContext,
    UserPrincipleEntry,
    VoiceModel,
)

if TYPE_CHECKING:
    from django.contrib.auth.models import AbstractUser

logger = logging.getLogger(__name__)


class ConstitutionService:
    """
    Composes dynamic constitution from multiple dimensions.

    This is the core service that agents call to get their context.
    It loads data from YAML (base), Django models (user), and performs
    semantic search on the second brain.

    Usage:
        service = ConstitutionService(user, "tenant_schema")
        constitution = service.compose(
            agent_type="v_accountant",
            query_context="Help me with quarterly taxes",
        )
        prompt = service.to_system_prompt(constitution)
    """

    def __init__(self, user: "AbstractUser", tenant_schema: str):
        """
        Initialize the constitution service.

        Args:
            user: The authenticated Django user
            tenant_schema: PostgreSQL schema name for multi-tenant queries
        """
        self.user = user
        self.tenant_schema = tenant_schema
        self._base_constitution: dict[str, Any] | None = None
        self._user_identity: UserIdentityContext | None = None
        self._user_principles: list[UserPrincipleEntry] | None = None

    def compose(
        self,
        agent_type: str,
        query_context: str = "",
        max_second_brain_entries: int = 10,
    ) -> ComposedConstitution:
        """
        Compose a complete constitution for this agent execution.

        Merges all 4 layers with proper priority handling:
        1. Load base constitution from YAML
        2. Load user identity from Django
        3. Load user principles filtered by agent type
        4. Query second brain with semantic search
        5. Merge with priority: User > Base

        Args:
            agent_type: The agent being invoked (v_accountant, v_creative, etc.)
            query_context: The user's query for semantic search
            max_second_brain_entries: How many knowledge entries to include

        Returns:
            ComposedConstitution with all layers merged
        """
        # 1. Load base constitution (cached)
        base_principles = ConstitutionConfig.get_principles()
        voice_models = ConstitutionConfig.get_voice_models()
        books = ConstitutionConfig.get_books()
        transformation = ConstitutionConfig.get_transformation_promise()
        philosophy = ConstitutionConfig.get_philosophy()
        absolute_constraints = ConstitutionConfig.get_absolute_constraints()
        soft_constraints = ConstitutionConfig.get_soft_constraints()
        agent_defaults = ConstitutionConfig.get_agent_defaults(agent_type)

        # 2. Load user identity
        identity = self._load_user_identity()

        # 3. Load user principles (filtered by agent)
        user_principles = self._load_user_principles(agent_type)

        # 4. Query second brain (semantic + weight-based)
        knowledge = self._query_second_brain(query_context, max_second_brain_entries)

        # 5. Merge principles with priority
        merged_principles = self._merge_principles(
            base_principles,
            user_principles,
            agent_defaults,
            agent_type,
        )

        # 6. Select voice model
        voice_model = self._select_voice_model(voice_models, agent_defaults)

        # 7. Filter books by relevance
        relevant_books = self._filter_books(books, agent_defaults, query_context)

        return ComposedConstitution(
            version=ConstitutionConfig.get_version(),
            agent_type=agent_type,
            principles=merged_principles,
            voice_model=voice_model,
            transformation_promise=transformation,
            philosophy=philosophy,
            identity=identity,
            user_principles=user_principles,
            knowledge=knowledge,
            reference_books=relevant_books,
            absolute_constraints=absolute_constraints,
            soft_constraints=soft_constraints,
        )

    def to_system_prompt(self, constitution: ComposedConstitution) -> str:
        """
        Convert composed constitution to system prompt text.

        This generates the actual text that will be included in the
        agent's system prompt.

        Args:
            constitution: The composed constitution

        Returns:
            Formatted system prompt string
        """
        return format_constitution_prompt(constitution)

    def _load_user_identity(self) -> UserIdentityContext | None:
        """
        Load user identity from Django model.

        Returns None if user has no identity profile.
        """
        if self._user_identity is not None:
            return self._user_identity

        try:
            # Try to get identity from user (requires Django model)
            identity = getattr(self.user, "identity", None)
            if identity is not None:
                context_dict = identity.to_constitution_context()
                self._user_identity = UserIdentityContext.model_validate(context_dict)
                return self._user_identity
        except Exception as e:
            logger.warning(f"Failed to load user identity: {e}")

        return None

    def _load_user_principles(self, agent_type: str) -> list[UserPrincipleEntry]:
        """
        Load user principles filtered by agent type.

        Only returns active principles that apply to this agent type.
        """
        if self._user_principles is not None:
            # Filter cached principles by agent type
            return [p for p in self._user_principles if not p.applies_to_agents or agent_type in p.applies_to_agents]

        try:
            # Try to get principles from user
            principles_qs = getattr(self.user, "principles", None)
            if principles_qs is not None:
                # Filter active principles
                db_principles = principles_qs.filter(is_active=True)

                # Convert to Pydantic models
                self._user_principles = []
                for p in db_principles:
                    context = p.to_constitution_context()
                    entry = UserPrincipleEntry.model_validate(context)
                    self._user_principles.append(entry)

                # Filter by agent type
                return [
                    p for p in self._user_principles if not p.applies_to_agents or agent_type in p.applies_to_agents
                ]
        except Exception as e:
            logger.warning(f"Failed to load user principles: {e}")

        return []

    def _query_second_brain(
        self,
        query: str,
        limit: int,
    ) -> list[SecondBrainEntryContext]:
        """
        Query second brain with semantic similarity + weight ranking.

        Merges both personal (tenant=None) and tenant-specific entries.
        Uses embedding service for semantic search.
        """
        if not query:
            return []

        try:
            from lightwave.ai.embeddings import EmbeddingService

            # Import Django models (optional dependency)
            try:
                from apps.core.knowledge.models import SecondBrainEntry
                from django.db.models import Q
            except ImportError:
                logger.debug("Second brain models not available")
                return []

            # Get query embedding
            embedding_service = EmbeddingService()
            query_result = embedding_service.embed("query", "temp", query)
            query_embedding = query_result.embedding

            # Find entries from BOTH layers:
            # 1. Personal (tenant=None) - follows user everywhere
            # 2. Tenant-specific - only in current tenant context
            entries = SecondBrainEntry.objects.filter(
                Q(user=self.user, tenant__isnull=True)  # Personal layer
                | Q(user=self.user, tenant__schema_name=self.tenant_schema),  # Tenant layer
                embedding__isnull=False,
            )

            # Score by similarity * weight
            scored: list[tuple[Any, float]] = []
            for entry in entries:
                if entry.embedding:
                    similarity = embedding_service.cosine_similarity(
                        query_embedding,
                        entry.embedding,
                    )
                    score = similarity * entry.effective_weight
                    scored.append((entry, score))

            # Sort by score, take top N
            scored.sort(key=lambda x: x[1], reverse=True)

            # Mark accessed and convert to Pydantic
            results = []
            for entry, _score in scored[:limit]:
                entry.mark_accessed()
                context = entry.to_constitution_context()
                results.append(SecondBrainEntryContext.model_validate(context))

            return results

        except Exception as e:
            logger.warning(f"Failed to query second brain: {e}")
            return []

    def _merge_principles(
        self,
        base_principles: list[Principle],
        user_principles: list[UserPrincipleEntry],
        agent_defaults: AgentDefaults,
        agent_type: str,
    ) -> list[Principle]:
        """
        Merge base and user principles with priority handling.

        1. Start with base principles
        2. Apply agent-specific boosts/reductions
        3. Apply user overrides (higher priority)
        4. Filter by applicability
        """
        # Create a dict of base principles by ID
        merged: dict[str, Principle] = {}
        for p in base_principles:
            # Check if this principle applies to this agent type
            if "all" in p.applies_to or agent_type in p.applies_to:
                merged[p.id] = p

        # Apply agent-specific boosts
        for boost_id in agent_defaults.principles_boost:
            if boost_id in merged:
                old_principle = merged[boost_id]
                merged[boost_id] = Principle(
                    id=old_principle.id,
                    name=old_principle.name,
                    description=old_principle.description,
                    weight=min(2.0, old_principle.weight * 1.3),  # 30% boost
                    category=old_principle.category,
                    applies_to=old_principle.applies_to,
                )

        # Apply agent-specific reductions
        for reduce_id in agent_defaults.principles_reduce:
            if reduce_id in merged:
                old_principle = merged[reduce_id]
                merged[reduce_id] = Principle(
                    id=old_principle.id,
                    name=old_principle.name,
                    description=old_principle.description,
                    weight=max(0.1, old_principle.weight * 0.5),  # 50% reduction
                    category=old_principle.category,
                    applies_to=old_principle.applies_to,
                )

        # Apply user overrides (highest priority)
        for user_p in user_principles:
            if user_p.overrides_base and user_p.overrides_base in merged:
                # Override the base principle
                merged[user_p.overrides_base] = Principle(
                    id=user_p.overrides_base,
                    name=user_p.name,
                    description=user_p.description,
                    weight=user_p.weight,
                    category=user_p.category,
                    applies_to=user_p.applies_to_agents or ["all"],
                )

        # Sort by weight descending
        sorted_principles = sorted(merged.values(), key=lambda p: p.weight, reverse=True)
        return sorted_principles

    def _select_voice_model(
        self,
        voice_models: list[VoiceModel],
        agent_defaults: AgentDefaults,
    ) -> VoiceModel | None:
        """Select the appropriate voice model for this agent."""
        target_id = agent_defaults.voice_model
        for vm in voice_models:
            if vm.id == target_id:
                return vm
        # Fallback to first voice model
        return voice_models[0] if voice_models else None

    def _filter_books(
        self,
        books: list[KnowledgeBook],
        agent_defaults: AgentDefaults,
        query_context: str,
    ) -> list[KnowledgeBook]:
        """Filter books by relevance to agent type and query."""
        if not agent_defaults.knowledge_filter:
            # No filter - return top books by weight
            return sorted(books, key=lambda b: b.weight, reverse=True)[:5]

        # Filter by relevance tags
        relevant = []
        for book in books:
            for tag in book.relevance:
                if any(f in tag.lower() for f in agent_defaults.knowledge_filter):
                    relevant.append(book)
                    break

        return sorted(relevant, key=lambda b: b.weight, reverse=True)[:5]


# =============================================================================
# CONVENIENCE FUNCTIONS
# =============================================================================


def compose_constitution(
    user: "AbstractUser",
    tenant_schema: str,
    agent_type: str,
    query_context: str = "",
    max_second_brain_entries: int = 10,
) -> ComposedConstitution:
    """
    Compose a constitution for an agent execution.

    Convenience function that creates a service and composes in one call.

    Args:
        user: The authenticated Django user
        tenant_schema: PostgreSQL schema name
        agent_type: The agent being invoked
        query_context: The user's query for semantic search
        max_second_brain_entries: How many knowledge entries to include

    Returns:
        ComposedConstitution with all layers merged
    """
    service = ConstitutionService(user, tenant_schema)
    return service.compose(
        agent_type=agent_type,
        query_context=query_context,
        max_second_brain_entries=max_second_brain_entries,
    )
