from typing import List, Dict, Any, Optional, Callable
import logging
import json
from sqlalchemy.orm import Session
from app.db import models

class ContextManager:
    """Handles prompt assembly and conversation history management WITHOUT regex."""
    
    def __init__(self, history_formatter: Optional[Callable] = None, context_postprocessor: Optional[Callable] = None):
        self.history_formatter = history_formatter or self._default_history_formatter
        self.context_postprocessor = context_postprocessor or self._default_context_postprocessor

    def prepare_initial_messages(self, question: str, context_chunks: List[Any], history: List[models.Message], 
                                feature_name: str, mesh_context: str, sync_workspace_id: Optional[str],
                                db: Optional[Session] = None, user_id: Optional[str] = None,
                                prompt_service: Any = None, prompt_slug: str = "rag-pipeline",
                                tools: List[Dict[str, Any]] = None,
                                session_override: Optional[str] = None) -> List[Dict[str, str]]:
        
        from .profiles import get_profile
        profile = get_profile(feature_name)
        
        history_text = self.history_formatter(history)
        context_text = self.context_postprocessor(context_chunks)

        template = profile.template
        
        # dynamic prompt override
        if prompt_service and db and user_id:
            db_prompt = prompt_service.get_prompt_by_slug(db, prompt_slug, user_id)
            if db_prompt:
                template = db_prompt.content
                
        # Session Specific Override takes maximum priority
        if session_override and session_override.strip():
            template = session_override
        
        # Augmented mesh info
        mesh_info = mesh_context
        attached_count = mesh_info.count("- Node ID:")
        
        # Factual infrastructure summary (Neutral context)
        inventory = f"### 🗺️ Infrastructure Inventory ({attached_count} nodes)\n"
        if attached_count == 0:
            inventory += "- No nodes are currently attached to this session.\n"
        else:
            inventory += mesh_info
            if "Status: offline" in mesh_info.lower():
                inventory += "\nNote: Certain attached nodes are currently offline and unreachable.\n"
        
        mesh_info = inventory

        if sync_workspace_id:
            mesh_info += f"\nActive Ghost Mirror Session ID: {sync_workspace_id}\n"
            mesh_info += f"Instruction: Your `session_id` is automatically injected into tools like `mesh_file_explorer` or `mesh_sync_control`. You do not need to provide it.\n"

        # List of actual tools for strict enforcement
        available_skills = [t["function"]["name"] for t in (tools or [])]
        skill_list_str = ", ".join(available_skills) if available_skills else "NONE"

        # Robust string replacement to avoid crashing on braces in context/history
        system_prompt = template.replace("{question}", str(question))
        system_prompt = system_prompt.replace("{context}", str(context_text))
        system_prompt = system_prompt.replace("{chat_history}", str(history_text))
        system_prompt = system_prompt.replace("{mesh_context}", str(mesh_info))
        
        # Handle escaped session_id marker in DEFAULT_PROMPT_TEMPLATE
        system_prompt = system_prompt.replace("{{session_id}}", "{session_id}")

        # Enforce Tool-belt awareness — only inject when tools actually exist
        if available_skills:
            system_prompt += f"\n\n## 🛠️ ACTIVE TOOL-BELT (FORBIDDEN to mention others):\n"
            system_prompt += f"You have ONLY these {len(available_skills)} calibrated tools: [{skill_list_str}].\n"
            system_prompt += "If a user asks about your capabilities, you MUST only list these specific tools. DO NOT promise Google Calendar, Wolfram Alpha, or any external integrations not in this list."

        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": question}
        ]

    async def compress_history(self, messages: List[Dict[str, Any]], llm_provider=None) -> List[Dict[str, Any]]:
        """Summarizes middle turns using the LLM to preserve context window."""
        if len(messages) < 60:
            return messages
        
        logging.info(f"[ContextManager] Memory overflow ({len(messages)}). Compressing middle history via LLM...")
        system_msg = messages[0]
        pivot_idx = len(messages) - 20
        while pivot_idx > 1 and messages[pivot_idx].get("role") != "user":
            pivot_idx -= 1
        
        if pivot_idx <= 1: pivot_idx = 10
            
        middle_to_compress = messages[1:pivot_idx]
        recent_msg = messages[pivot_idx:]
        
        summary_lines = []
        for m in middle_to_compress:
            role = m.get("role", "unknown")
            content = m.get("content") or ""
            if role == "user":
                summary_lines.append(f"User: {content[:100]}")
            elif role == "assistant" and not m.get("tool_calls"):
                summary_lines.append(f"AI: {content[:100]}")
            elif role == "tool":
                summary_lines.append(f"Result: {content[:100]}")
        
        summary_text = "\n".join(summary_lines)
        
        if llm_provider:
             try:
                 prompt = f"Please compress the following timeline of chronological actions into a dense, factual summary paragraph. Omit fluff.\n\n{summary_text}"
                 comp_res = await llm_provider.acompletion(
                     messages=[{"role": "user", "content": prompt}],
                     max_tokens=500,
                     temperature=0.0,
                     stream=False
                 )
                 if comp_res and comp_res.choices:
                      try:
                           llm_summary = comp_res.choices[0].message.content
                      except AttributeError:
                           llm_summary = comp_res.choices[0].get("message", {}).get("content")
                      if llm_summary:
                           summary_text = llm_summary.strip()
             except Exception as e:
                 logging.warning(f"LLM Summary compression failed: {e}")
        
        history_msg = {
            "role": "user", 
            "content": f"[SYSTEM: Historical Context Timeline Condensed into Summary]\n\n{summary_text}\n\n[End Summary.]"
        }
        ack_msg = {
            "role": "assistant",
            "content": "Understood. Proceeding with current state."
        }
        
        return [system_msg, history_msg, ack_msg] + recent_msg

    def _default_context_postprocessor(self, contexts: List[Any]) -> str:
        """Safe context joining that handles both raw strings and enriched dictionaries."""
        processed = []
        for c in contexts:
            if isinstance(c, str):
                processed.append(c)
            elif isinstance(c, dict):
                # Try common keys for content
                content = c.get("content") or c.get("text") or str(c)
                processed.append(content)
            else:
                processed.append(str(c))
        return "\n\n".join(processed) or "No context provided."

    def _default_history_formatter(self, history: List[models.Message]) -> str:
        """High-performance history formatting using string methods instead of regex."""
        lines = []
        for msg in history:
            role = "Human" if msg.sender == "user" else "Assistant"
            content = msg.content or ""
            
            if role == "Assistant":
                if not content:
                    content = "[Action: Calling tools or internal reasoning...]"
                else:
                    # Strip turn headers without regex
                    content_lines = content.splitlines(keepends=True)
                    cleaned_content_lines = []
                    for line in content_lines:
                        l_lower = line.lower().strip()
                        if "---" in l_lower and ("turn" in l_lower or "analysis" in l_lower):
                             continue
                        if l_lower.startswith("###") and ("analysis" in l_lower or "turn" in l_lower):
                             continue
                        cleaned_content_lines.append(line)
                    content = "".join(cleaned_content_lines).strip()
            
            lines.append(f"{role}: {content.strip()}")
        return "\n".join(lines)
