"""
Database-Driven Constitution Service with Token Budgeting.

Enhanced ConstitutionService that:
1. Loads constitution items from database (prompts, rules, references)
2. Token-budgets context assembly by priority
3. Integrates test breadcrumbs for session continuity

SST Source: flows/automation/agentic/enrichment/docs-alignment.yaml

Token Budget Priority:
    1000 - Absolute constraints (always included)
     900 - Active task context (25% budget)
     800 - Relevant rules (20% budget)
     700 - Test breadcrumbs (15% budget)
     600 - Second brain knowledge (25% budget)
     500 - Reference documents (15% budget)

Usage:
    from lightwave.ai.constitution.database_service import DatabaseConstitutionService

    service = DatabaseConstitutionService(site, max_tokens=100000)
    context = service.assemble(
        agent_type="software-architect",
        task_id="abc123",
        file_paths=["apps/users/views.py"],
    )
    prompt = service.to_system_prompt(context)
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

# Use structlog if available, fallback to standard logging
try:
    import structlog

    logger = structlog.get_logger(__name__)
except ImportError:
    logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    pass


# =============================================================================
# CONSTANTS
# =============================================================================

# Token budget shares by priority tier
BUDGET_ALLOCATION = {
    1000: 0.0,  # Absolute constraints - always full
    900: 0.25,  # Active task context
    800: 0.20,  # Relevant rules
    700: 0.15,  # Test breadcrumbs
    600: 0.25,  # Second brain knowledge
    500: 0.15,  # Reference documents
}

# Default token budgets by item type
DEFAULT_TOKEN_LIMITS = {
    "prompt": 5000,
    "rule": 2000,
    "reference": 3000,
    "test_breadcrumb": 1500,
    "task_context": 4000,
}

# Categories that are part of Constitution
CONSTITUTION_CATEGORIES = [
    "prompt",
    "rule",
    "reference",
    "agent_definition",
    "skill_definition",
    "test_breadcrumb",
]


# =============================================================================
# DATA CLASSES
# =============================================================================


@dataclass
class ConstitutionItem:
    """A single item to include in constitution context."""

    id: str
    category: str
    title: str
    content: str
    priority: int
    token_count: int
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class TokenBudget:
    """Token budget tracking for context assembly."""

    total_budget: int
    allocated: dict[int, int] = field(default_factory=dict)
    used: dict[int, int] = field(default_factory=dict)

    def __post_init__(self) -> None:
        # Calculate allocations based on BUDGET_ALLOCATION
        remaining = self.total_budget
        for priority, share in sorted(BUDGET_ALLOCATION.items(), reverse=True):
            if share == 0.0:
                self.allocated[priority] = remaining  # Absolute gets everything
            else:
                allocation = int(self.total_budget * share)
                self.allocated[priority] = allocation
                remaining -= allocation
            self.used[priority] = 0

    def can_fit(self, priority: int, tokens: int) -> bool:
        """Check if item fits in budget for this priority."""
        available = self.allocated.get(priority, 0) - self.used.get(priority, 0)
        return tokens <= available

    def consume(self, priority: int, tokens: int) -> bool:
        """Consume tokens from a priority tier. Returns success."""
        if not self.can_fit(priority, tokens):
            return False
        self.used[priority] = self.used.get(priority, 0) + tokens
        return True

    def remaining(self, priority: int) -> int:
        """Get remaining budget for priority tier."""
        return self.allocated.get(priority, 0) - self.used.get(priority, 0)


@dataclass
class AssembledContext:
    """Result of context assembly with token budgeting."""

    agent_type: str
    items: list[ConstitutionItem] = field(default_factory=list)
    budget: TokenBudget | None = None
    task_context: dict[str, Any] | None = None
    breadcrumbs: list[dict[str, Any]] = field(default_factory=list)

    @property
    def total_tokens(self) -> int:
        """Total tokens used."""
        return sum(item.token_count for item in self.items)


# =============================================================================
# DATABASE CONSTITUTION SERVICE
# =============================================================================


class DatabaseConstitutionService:
    """
    Database-driven constitution service with token budgeting.

    Loads constitution items (prompts, rules, references) from database
    and assembles them into a token-budgeted context for Claude sessions.
    """

    def __init__(
        self,
        site: Any = None,
        max_tokens: int = 100000,
    ) -> None:
        """
        Initialize the database constitution service.

        Args:
            site: Site instance for database queries.
            max_tokens: Maximum tokens for assembled context.
        """
        self._site = site
        self.max_tokens = max_tokens

    @property
    def site(self) -> Any:
        """Get or resolve the site for queries."""
        if self._site is not None:
            return self._site

        try:
            from apps.platform.cms.models import Site

            self._site = Site.objects.first()
            if self._site is None:
                raise RuntimeError("No Site found")
            return self._site
        except ImportError:
            raise RuntimeError("Django not available")

    def assemble(
        self,
        agent_type: str,
        task_id: str | None = None,
        file_paths: list[str] | None = None,
        query_context: str = "",
    ) -> AssembledContext:
        """
        Assemble token-budgeted context for a Claude session.

        Args:
            agent_type: The agent type (e.g., "software-architect").
            task_id: Optional task ID for context.
            file_paths: Optional file paths to find relevant breadcrumbs.
            query_context: Optional query for semantic search.

        Returns:
            AssembledContext with prioritized items.
        """
        budget = TokenBudget(total_budget=self.max_tokens)
        items: list[ConstitutionItem] = []

        # 1. Load absolute constraints (priority 1000, always included)
        constraints = self._load_absolute_constraints(agent_type)
        for item in constraints:
            if budget.consume(1000, item.token_count):
                items.append(item)

        # 2. Load task context (priority 900)
        task_context = None
        if task_id:
            task_items, task_context = self._load_task_context(task_id)
            for item in task_items:
                if budget.consume(900, item.token_count):
                    items.append(item)

        # 3. Load rules (priority 800)
        rules = self._load_rules(agent_type)
        for item in sorted(rules, key=lambda x: x.priority, reverse=True):
            if budget.can_fit(800, item.token_count):
                budget.consume(800, item.token_count)
                items.append(item)

        # 4. Load breadcrumbs (priority 700)
        breadcrumbs = self._load_breadcrumbs(file_paths)
        for bc in breadcrumbs:
            item = self._breadcrumb_to_item(bc)
            if budget.can_fit(700, item.token_count):
                budget.consume(700, item.token_count)
                items.append(item)

        # 5. Load second brain knowledge (priority 600)
        if query_context:
            knowledge = self._load_knowledge(query_context)
            for item in knowledge:
                if budget.can_fit(600, item.token_count):
                    budget.consume(600, item.token_count)
                    items.append(item)

        # 6. Load references (priority 500)
        references = self._load_references(agent_type, query_context)
        for item in references:
            if budget.can_fit(500, item.token_count):
                budget.consume(500, item.token_count)
                items.append(item)

        return AssembledContext(
            agent_type=agent_type,
            items=items,
            budget=budget,
            task_context=task_context,
            breadcrumbs=[self._format_breadcrumb(bc) for bc in breadcrumbs],
        )

    def to_system_prompt(self, context: AssembledContext) -> str:
        """
        Convert assembled context to system prompt text.

        Args:
            context: The assembled context.

        Returns:
            Formatted system prompt string.
        """
        sections: list[str] = []

        # Header
        sections.append("# Constitution Context")
        sections.append(f"Agent: {context.agent_type}")
        sections.append(f"Token budget: {context.total_tokens}/{self.max_tokens}")
        sections.append("")

        # Group items by category
        by_category: dict[str, list[ConstitutionItem]] = {}
        for item in context.items:
            by_category.setdefault(item.category, []).append(item)

        # Format each category
        category_order = [
            ("prompt", "## System Prompts"),
            ("rule", "## Behavioral Rules"),
            ("test_breadcrumb", "## Test Context (Breadcrumbs)"),
            ("reference", "## Reference Material"),
            ("task", "## Task Context"),
        ]

        for category, header in category_order:
            items = by_category.get(category, [])
            if items:
                sections.append(header)
                sections.append("")
                for item in items:
                    sections.append(f"### {item.title}")
                    sections.append(item.content)
                    sections.append("")

        return "\n".join(sections)

    # =========================================================================
    # PRIVATE LOADING METHODS
    # =========================================================================

    def _load_absolute_constraints(self, agent_type: str) -> list[ConstitutionItem]:
        """Load absolute constraints (always included)."""
        try:
            from apps.platform.createos.document.models import Document

            docs = Document.objects.filter(
                site=self.site,
                category="prompt",
                is_deleted=False,
                status="published",
                metadata__prompt_type="system",
                metadata__priority__gte=900,
            ).order_by("-metadata__priority")[:5]

            return [self._document_to_item(doc, 1000) for doc in docs]

        except Exception as e:
            logger.warning("constraints_load_failed", error=str(e), agent_type=agent_type)
            return []

    def _load_task_context(
        self,
        task_id: str,
    ) -> tuple[list[ConstitutionItem], dict[str, Any] | None]:
        """Load context for a specific task."""
        try:
            from apps.platform.createos.task.models import Task

            task = Task.objects.get(id=task_id)

            # Build task context
            context = {
                "id": str(task.id),
                "title": task.title,
                "description": task.description or "",
                "status": task.status,
                "priority": task.priority,
            }

            # Create task item
            content = f"**{task.title}**\n\n{task.description or 'No description'}"
            item = ConstitutionItem(
                id=str(task.id),
                category="task",
                title=f"Task: {task.title}",
                content=content,
                priority=900,
                token_count=self._estimate_tokens(content),
                metadata=context,
            )

            return [item], context

        except Exception as e:
            logger.warning("task_context_load_failed", error=str(e), task_id=task_id)
            return [], None

    def _load_rules(self, agent_type: str) -> list[ConstitutionItem]:
        """Load rules applicable to this agent type."""
        try:
            from apps.platform.createos.document.models import Document
            from django.db.models import Q

            # Query rules that apply to this agent or all agents
            docs = (
                Document.objects.filter(
                    site=self.site,
                    category="rule",
                    is_deleted=False,
                    status="published",
                )
                .filter(Q(metadata__applies_to_agents=[]) | Q(metadata__applies_to_agents__contains=[agent_type]))
                .order_by("-metadata__priority")[:10]
            )

            return [self._document_to_item(doc, doc.metadata.get("priority", 100)) for doc in docs]

        except Exception as e:
            logger.warning("rules_load_failed", error=str(e), agent_type=agent_type)
            return []

    def _load_breadcrumbs(
        self,
        file_paths: list[str] | None,
    ) -> list[Any]:
        """Load breadcrumbs for given file paths."""
        try:
            from apps.platform.createos.claude.services.breadcrumbs import (
                BreadcrumbService,
            )

            service = BreadcrumbService(site=self.site)
            return list(
                service.get_context_for_session(
                    file_paths=file_paths,
                    max_breadcrumbs=5,
                )
            )

        except Exception as e:
            logger.warning("breadcrumbs_load_failed", error=str(e))
            return []

    def _load_knowledge(self, query_context: str) -> list[ConstitutionItem]:
        """Load second brain knowledge via semantic search."""
        # This would integrate with the existing second brain
        # For now, return empty - can be extended
        return []

    def _load_references(
        self,
        agent_type: str,
        query_context: str,
    ) -> list[ConstitutionItem]:
        """Load reference documents."""
        try:
            from apps.platform.createos.document.models import Document

            docs = Document.objects.filter(
                site=self.site,
                category="reference",
                is_deleted=False,
                status="published",
            ).order_by("-metadata__priority")[:5]

            return [self._document_to_item(doc, doc.metadata.get("priority", 50)) for doc in docs]

        except Exception as e:
            logger.warning("references_load_failed", error=str(e), agent_type=agent_type)
            return []

    # =========================================================================
    # HELPER METHODS
    # =========================================================================

    def _document_to_item(
        self,
        doc: Any,
        priority: int,
    ) -> ConstitutionItem:
        """Convert a Document to a ConstitutionItem."""
        # Extract content
        if "raw_markdown" in doc.content:
            content = doc.content["raw_markdown"]
        elif "blocks" in doc.content:
            content = "\n\n".join(b.get("content", "") for b in doc.content.get("blocks", []))
        else:
            content = str(doc.content)

        return ConstitutionItem(
            id=str(doc.id),
            category=doc.category,
            title=doc.metadata.get("title", doc.title or "Untitled"),
            content=content,
            priority=priority,
            token_count=self._estimate_tokens(content),
            metadata=doc.metadata,
        )

    def _breadcrumb_to_item(self, breadcrumb: dict[str, Any]) -> ConstitutionItem:
        """Convert a breadcrumb dict to a ConstitutionItem."""
        content = f"""**Test Failure: {breadcrumb.get('test_function', 'Unknown')}**
File: {breadcrumb.get('test_file', 'Unknown')}
Type: {breadcrumb.get('failure_type', 'assertion')}

Error:
{breadcrumb.get('error_message', 'No error message')}

Suggested fixes:
{chr(10).join('- ' + fix for fix in breadcrumb.get('suggested_fixes', []))}

Previously tried:
{chr(10).join('- ' + tried for tried in breadcrumb.get('what_was_tried', []))}
"""
        return ConstitutionItem(
            id=breadcrumb.get("id", ""),
            category="test_breadcrumb",
            title=f"Breadcrumb: {breadcrumb.get('test_function', 'Unknown')}",
            content=content,
            priority=700,
            token_count=self._estimate_tokens(content),
            metadata=breadcrumb,
        )

    def _format_breadcrumb(self, breadcrumb: dict[str, Any]) -> dict[str, Any]:
        """Format breadcrumb for context output."""
        return {
            "id": breadcrumb.get("id"),
            "test_function": breadcrumb.get("test_function"),
            "error_summary": breadcrumb.get("error_message", "")[:100],
        }

    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count for text (rough approximation)."""
        # Rough estimate: 1 token ≈ 4 characters
        return len(text) // 4


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


def assemble_constitution(
    agent_type: str,
    task_id: str | None = None,
    file_paths: list[str] | None = None,
    query_context: str = "",
    max_tokens: int = 100000,
) -> AssembledContext:
    """
    Assemble a token-budgeted constitution for a Claude session.

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

    Args:
        agent_type: The agent type.
        task_id: Optional task ID for context.
        file_paths: Optional file paths for breadcrumb lookup.
        query_context: Optional query for semantic search.
        max_tokens: Maximum tokens for context.

    Returns:
        AssembledContext with prioritized items.
    """
    service = DatabaseConstitutionService(max_tokens=max_tokens)
    return service.assemble(
        agent_type=agent_type,
        task_id=task_id,
        file_paths=file_paths,
        query_context=query_context,
    )
