"""Context management: compaction, token estimation, and error tracking."""

import time

from genworker.config import MAX_CONTEXT_IMAGES, log


def compact_context(messages: list, max_images: int = MAX_CONTEXT_IMAGES) -> list:
    """Remove old screenshot data from messages to prevent context overflow.

    Replaces image blocks (oldest first) with a placeholder text so the
    action history text is still visible but image tokens are freed.
    """
    image_count = 0
    for msg in messages:
        if not isinstance(msg.get("content"), list):
            continue
        for item in msg["content"]:
            if not isinstance(item, dict):
                continue
            if item.get("type") == "image":
                image_count += 1
            elif item.get("type") == "tool_result":
                for sub in item.get("content", []):
                    if isinstance(sub, dict) and sub.get("type") == "image":
                        image_count += 1

    if image_count <= max_images:
        return messages

    to_remove = image_count - max_images
    removed   = 0

    for msg in messages:
        if removed >= to_remove:
            break
        if not isinstance(msg.get("content"), list):
            continue
        new_content = []
        for item in msg["content"]:
            if isinstance(item, dict) and item.get("type") == "image" and removed < to_remove:
                new_content.append({"type": "text", "text": "[screenshot removed — old context]"})
                removed += 1
            elif isinstance(item, dict) and item.get("type") == "tool_result":
                new_sub = []
                for sub in item.get("content", []):
                    if isinstance(sub, dict) and sub.get("type") == "image" and removed < to_remove:
                        new_sub.append({"type": "text", "text": "[screenshot removed]"})
                        removed += 1
                    else:
                        new_sub.append(sub)
                new_content.append({**item, "content": new_sub})
            else:
                new_content.append(item)
        msg["content"] = new_content

    log.info(f"Context compaction: removed {removed} old screenshots (had {image_count}, target {max_images})")
    return messages


def estimate_context_size(messages: list) -> dict:
    """Rough estimate of token usage in the current message list.

    Uses ~4 chars/token for text and ~1600 tokens per JPEG screenshot.
    """
    text_chars  = 0
    image_count = 0

    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            text_chars += len(content)
        elif isinstance(content, list):
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        text_chars += len(item.get("text", ""))
                    elif item.get("type") == "image":
                        image_count += 1
                    elif item.get("type") == "tool_result":
                        for sub in item.get("content", []):
                            if isinstance(sub, dict):
                                if sub.get("type") == "text":
                                    text_chars += len(sub.get("text", ""))
                                elif sub.get("type") == "image":
                                    image_count += 1
                elif hasattr(item, "type") and item.type == "text":
                    text_chars += len(getattr(item, "text", ""))

    est_tokens = (text_chars // 4) + (image_count * 1600)
    return {"text_chars": text_chars, "image_count": image_count, "est_tokens": est_tokens}


class ErrorTracker:
    """Track consecutive failures and detect stuck-loops to trigger recovery."""

    def __init__(self):
        self.consecutive_errors = 0
        self.error_history: list[dict] = []
        self.last_actions:  list[str]  = []

    def record_success(self, action: str):
        self.consecutive_errors = 0
        self.last_actions.append(action)
        if len(self.last_actions) > 10:
            self.last_actions = self.last_actions[-10:]

    def record_error(self, action: str, error: str):
        self.consecutive_errors += 1
        self.error_history.append({"action": action, "error": error, "time": time.time()})
        self.last_actions.append(f"ERROR:{action}")
        if len(self.error_history) > 20:
            self.error_history = self.error_history[-20:]

    def needs_recovery(self) -> bool:
        return self.consecutive_errors >= 2

    def get_recovery_hint(self) -> str:
        recent  = self.error_history[-3:] if self.error_history else []
        actions = [e["action"] for e in recent]
        errors  = [e["error"]  for e in recent]
        return (
            f"⚠️ RECOVERY NEEDED — {self.consecutive_errors} consecutive errors detected.\n"
            f"Recent failed actions: {actions}\n"
            f"Recent errors: {errors}\n"
            f"SUGGESTIONS:\n"
            f"- Take a screenshot to re-orient yourself\n"
            f"- Try a completely different approach\n"
            f"- If GUI interaction keeps failing, try bash/CLI instead\n"
            f"- If a file operation fails, verify the path exists first\n"
            f"- If clicking fails, the UI may have changed — screenshot first\n"
        )

    def is_stuck_in_loop(self) -> bool:
        """Detect when the agent keeps repeating the same 1–2 actions."""
        if len(self.last_actions) < 4:
            return False
        return len(set(self.last_actions[-4:])) <= 2
