"""
Brynq Context Window Optimizer (Phase 64)
Helps agents work within token limits by summarizing, chunking, and compressing
context before sending it to an LLM.

Storage:
  state/broker/context_cache/  — cached summaries
"""

import json
import math
import re
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from .server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

CONTEXT_CACHE_DIR = BROKER_DIR / "context_cache"


def _ensure_dirs():
    CONTEXT_CACHE_DIR.mkdir(parents=True, exist_ok=True)


def estimate_tokens(text: str) -> int:
    """Rough token count estimation: words * 1.3."""
    if not text or not text.strip():
        return 0
    words = len(text.split())
    return int(math.ceil(words * 1.3))


def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
    """Split text into overlapping chunks for processing.

    Args:
        text: The text to split.
        chunk_size: Max number of characters per chunk.
        overlap: Number of characters to overlap between chunks.

    Returns:
        List of text chunks.
    """
    if not text:
        return []
    if len(text) <= chunk_size:
        return [text]

    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap
        if start >= len(text):
            break
    return chunks


def extract_key_points(messages: List[Dict[str, Any]]) -> Dict[str, List[str]]:
    """Pull out decisions, action items, questions, and status updates from messages.

    Scans message content for patterns indicating each category.

    Returns:
        Dict with keys: decisions, action_items, questions, status_updates.
    """
    decisions: List[str] = []
    action_items: List[str] = []
    questions: List[str] = []
    status_updates: List[str] = []

    decision_patterns = [
        r"(?i)\bdecided\b",
        r"(?i)\bagreed\b",
        r"(?i)\bwill use\b",
        r"(?i)\bgoing with\b",
        r"(?i)\blet'?s go with\b",
        r"(?i)\bapproved\b",
        r"(?i)\bdecision:\s*",
    ]
    action_patterns = [
        r"(?i)\btodo\b",
        r"(?i)\baction item\b",
        r"(?i)\bneed to\b",
        r"(?i)\bshould\b",
        r"(?i)\bwill do\b",
        r"(?i)\bassigned to\b",
        r"(?i)\btask:\s*",
    ]
    question_patterns = [
        r"\?$",
        r"\?\s",
        r"(?i)\bquestion:\s*",
        r"(?i)\banyone know\b",
        r"(?i)\bwhat about\b",
    ]
    status_patterns = [
        r"(?i)\bcompleted?\b",
        r"(?i)\bdone\b",
        r"(?i)\bfinished\b",
        r"(?i)\bin progress\b",
        r"(?i)\bblocked\b",
        r"(?i)\bstarted\b",
        r"(?i)\bstatus:\s*",
        r"(?i)\bupdate:\s*",
    ]

    for msg in messages:
        content = msg.get("message", msg.get("content", ""))
        if not content:
            continue

        for pattern in decision_patterns:
            if re.search(pattern, content):
                decisions.append(content.strip())
                break

        for pattern in action_patterns:
            if re.search(pattern, content):
                action_items.append(content.strip())
                break

        for pattern in question_patterns:
            if re.search(pattern, content):
                questions.append(content.strip())
                break

        for pattern in status_patterns:
            if re.search(pattern, content):
                status_updates.append(content.strip())
                break

    return {
        "decisions": decisions,
        "action_items": action_items,
        "questions": questions,
        "status_updates": status_updates,
    }


def summarize_conversation(messages: List[Dict[str, Any]],
                           max_tokens: int = 1000) -> str:
    """Summarize a list of messages into a concise context string.

    Groups key points and truncates to fit within max_tokens.
    """
    if not messages:
        return ""

    key_points = extract_key_points(messages)

    parts = []
    if key_points["decisions"]:
        parts.append("Decisions: " + "; ".join(key_points["decisions"][:5]))
    if key_points["action_items"]:
        parts.append("Action items: " + "; ".join(key_points["action_items"][:5]))
    if key_points["questions"]:
        parts.append("Open questions: " + "; ".join(key_points["questions"][:3]))
    if key_points["status_updates"]:
        parts.append("Status: " + "; ".join(key_points["status_updates"][:5]))

    # If no key points extracted, fall back to recent messages
    if not parts:
        recent = messages[-10:]
        for msg in recent:
            sender = msg.get("sender", msg.get("session_name", "unknown"))
            content = msg.get("message", msg.get("content", ""))
            if content:
                parts.append(f"{sender}: {content[:200]}")

    summary = "\n".join(parts)

    # Truncate to fit max_tokens
    while estimate_tokens(summary) > max_tokens and len(summary) > 100:
        summary = summary[:int(len(summary) * 0.8)]

    return summary.strip()


def summarize_channel(channel: str, max_tokens: int = 2000) -> str:
    """Summarize a channel's history into a concise context string.

    Reads the channel JSONL file and produces a summary grouped by topic,
    keeping key decisions and action items.
    """
    _ensure_dirs()

    channel_file = BROKER_DIR / "channels" / f"{channel}.jsonl"
    messages = _read_jsonl_sync(channel_file)

    if not messages:
        return f"Channel '{channel}' has no messages."

    summary = summarize_conversation(messages, max_tokens=max_tokens)

    # Cache the summary
    cache_file = CONTEXT_CACHE_DIR / f"channel_{channel}.json"
    _write_json_sync(cache_file, {
        "channel": channel,
        "message_count": len(messages),
        "summary": summary,
        "cached_at": datetime.now(timezone.utc).isoformat(),
    })

    return summary


def compress_context(context_dict: Dict[str, Any],
                     target_tokens: int = 4000) -> Dict[str, Any]:
    """Compress a context dict by removing low-priority fields and truncating long values.

    Priority order (highest to lowest): task, decisions, status, history, metadata.
    """
    if not context_dict:
        return {}

    # Priority levels for fields
    priority = {
        "task": 1,
        "task_description": 1,
        "task_id": 1,
        "decisions": 2,
        "action_items": 2,
        "status": 3,
        "status_updates": 3,
        "recent_activity": 4,
        "history": 5,
        "channel_summary": 4,
        "team_info": 4,
        "metadata": 6,
        "context": 5,
    }

    result = dict(context_dict)
    current_tokens = estimate_tokens(json.dumps(result, default=str))

    if current_tokens <= target_tokens:
        return result

    # Remove fields from lowest priority first
    sorted_fields = sorted(
        result.keys(),
        key=lambda k: priority.get(k, 10),
        reverse=True,
    )

    for field in sorted_fields:
        if estimate_tokens(json.dumps(result, default=str)) <= target_tokens:
            break

        value = result[field]
        if isinstance(value, str) and len(value) > 200:
            result[field] = value[:200] + "..."
        elif isinstance(value, list) and len(value) > 3:
            result[field] = value[:3]
        elif isinstance(value, dict) and len(json.dumps(value, default=str)) > 300:
            # Truncate nested dict values
            truncated = {}
            for k, v in list(value.items())[:5]:
                if isinstance(v, str) and len(v) > 100:
                    truncated[k] = v[:100] + "..."
                else:
                    truncated[k] = v
            result[field] = truncated

    # If still too large, drop lowest-priority fields entirely
    for field in sorted_fields:
        if estimate_tokens(json.dumps(result, default=str)) <= target_tokens:
            break
        if priority.get(field, 10) >= 5:
            del result[field]

    return result


def build_agent_context(agent_id: str, task_description: str,
                        max_tokens: int = 4000) -> Dict[str, Any]:
    """Build an optimized context for an agent about to work on a task.

    Assembles relevant knowledge, recent channel activity, task details,
    and team info, then compresses to fit within max_tokens.
    """
    _ensure_dirs()

    context: Dict[str, Any] = {
        "agent_id": agent_id,
        "task_description": task_description,
        "built_at": datetime.now(timezone.utc).isoformat(),
    }

    # Load agent profile if exists
    profile_path = BROKER_DIR / "profiles" / f"{agent_id}.json"
    if profile_path.exists():
        try:
            profile = json.loads(profile_path.read_text("utf-8"))
            context["agent_profile"] = {
                "name": profile.get("name", agent_id),
                "capabilities": profile.get("capabilities", []),
                "specialization": profile.get("specialization", ""),
            }
        except (json.JSONDecodeError, OSError):
            pass

    # Load agent memory if exists
    memory_path = BROKER_DIR / "memory" / f"{agent_id}.jsonl"
    if memory_path.exists():
        memories = _read_jsonl_sync(memory_path)
        if memories:
            # Take most recent memories
            context["recent_knowledge"] = memories[-10:]

    # Load recent general channel activity
    general_channel = BROKER_DIR / "channels" / "general.jsonl"
    if general_channel.exists():
        messages = _read_jsonl_sync(general_channel)
        if messages:
            recent = messages[-20:]
            context["channel_summary"] = summarize_conversation(
                recent, max_tokens=max(500, max_tokens // 4)
            )

    # Check context cache for any relevant cached summaries
    if CONTEXT_CACHE_DIR.exists():
        for cache_file in CONTEXT_CACHE_DIR.iterdir():
            if cache_file.suffix == ".json":
                try:
                    cached = json.loads(cache_file.read_text("utf-8"))
                    if cached.get("summary"):
                        context.setdefault("cached_summaries", []).append({
                            "source": cache_file.stem,
                            "summary": cached["summary"][:200],
                        })
                except (json.JSONDecodeError, OSError):
                    continue

    # Compress to target
    context = compress_context(context, target_tokens=max_tokens)
    return context
