import json
import json5
import re
import html
import logging
from typing import List, Dict, Any, Optional, Tuple

logger = logging.getLogger(__name__)

# ============================================================
# Text Processing
# ============================================================
def extract_thinking(text: str) -> Tuple[Optional[str], str]:
    reasoning_content = None
    content = text

    if '<think>' in content and '</think>' in content:
        match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
        if match:
            reasoning_content = match.group(1).strip()
            content = content.replace(match.group(0), "").strip()
    elif '</think>' in content:
        match = re.search(r'^(.*?)</think>', content, re.DOTALL)
        if match:
            reasoning_content = match.group(1).strip()
            content = content.replace(match.group(0), "").strip()

    return reasoning_content, content

def parse_tool_call(text: str) -> Tuple[Optional[Dict], str]:
    tool_call_text = None
    content = text

    if '<tool_call>' in content and '</tool_call>' in content:
        match = re.search(r'<tool_call>(.*?)</tool_call>', content, re.DOTALL)
        if match:
            tool_call_text = match.group(1).strip()
            content = content.replace(match.group(0), "").strip()
    elif '</tool_call>' in content:
        match = re.search(r'^(.*?)</tool_call>', content, re.DOTALL)
        if match:
            tool_call_text = match.group(1).strip()
            content = content.replace(match.group(0), "").strip()

    if tool_call_text:
        try:
            if "```json" in tool_call_text:
                tool_call_text = tool_call_text.split("```json")[1].split("```")[0].strip()
            elif "```" in tool_call_text:
                tool_call_text = tool_call_text.split("```")[1].split("```")[0].strip()

            parsed = json5.loads(tool_call_text)
            return parsed, content
        except:
            pass

        func_match = re.search(r'<function=([\w.]+)>', tool_call_text)
        if func_match:
            tool_name = func_match.group(1)
            tool_args = {}
            params = re.finditer(r'<parameter=([\w]+)>\s*(.*?)\s*</parameter>', tool_call_text, re.DOTALL)
            for p in params:
                param_name = p.group(1)
                param_value = p.group(2).strip()
                if param_value.startswith('"') and param_value.endswith('"'):
                    param_value = param_value[1:-1]
                try:
                    if param_value.isdigit():
                        param_value = int(param_value)
                except:
                    pass
                tool_args[param_name] = param_value

            return {"name": tool_name, "arguments": tool_args}, content

    return None, content

def is_final_answer(text: str) -> bool:
    t = text.lower()
    return (
        ('<answer>' in t and '</answer>' in t) or
        'final answer:' in t or
        ('exact answer:' in t and 'confidence:' in t)
    )

def strip_goal_met_tag(text: str) -> str:
    return re.sub(r'<goal_met>\s*(?:true|false)\s*</goal_met>', '', text, flags=re.IGNORECASE).strip()

def extract_goal_met(text: str) -> Optional[bool]:
    match = re.search(r'<goal_met>\s*(true|false)\s*</goal_met>', text, re.IGNORECASE)
    if match:
        return match.group(1).lower() == "true"
    return None

# ============================================================
# HTML Rendering Helpers
# ============================================================
def render_user_message(question: str) -> str:
    escaped = html.escape(question)
    return f'''<div class="user-message-bubble">
<div class="user-message-content">{escaped}</div>
</div>'''

def render_round_badge(round_num: int, max_rounds: int) -> str:
    return f'<div class="round-badge">Round {round_num}/{max_rounds}</div>'

def render_thinking_collapsed(text: str) -> str:
    escaped = html.escape(text)
    preview = text[:100] + "..." if len(text) > 100 else text
    preview_escaped = html.escape(preview)
    return f'''<details class="thinking-collapsed">
<summary>Thought process: "{preview_escaped}"</summary>
<div class="thinking-content">{escaped}</div>
</details>'''

def render_tool_call(fn_name: str, args: dict, browser = None) -> str:
    border_colors = {
        "browser.search": "#667eea",
        "browser.open": "#4facfe",
        "browser.find": "#fa709a",
        "browser.navigate": "#667eea",
        "browser.snapshot": "#4facfe",
        "browser.click": "#fa709a",
        "browser.type": "#86efac",
        "browser.press": "#facc15"
    }
    border_color = border_colors.get(fn_name, "#9ca3af")

    return f'''<div class="tool-call-card" style="border-left: 3px solid {border_color};">
<div class="tool-info">
<div class="tool-name">{html.escape(str(fn_name))}</div>
<div class="tool-detail">{html.escape(json.dumps(args))}</div>
</div>
</div>'''

def render_tool_result(result: str, fn_name: str) -> str:
    border_color = "#9ca3af"
    formatted_result = html.escape(result)
    formatted_result = formatted_result.replace('\n', '<br>')

    return f'''<div class="result-card-expanded" style="border-left: 3px solid {border_color};">
<div class="result-header-expanded">Result: {html.escape(fn_name)}</div>
<div class="result-content-expanded" style="font-family: monospace; white-space: pre-wrap;">{formatted_result}</div>
</div>'''

def render_completion() -> str:
    return '<div class="completion-msg">Task Complete</div>'

# ============================================================
# Token Management
# ============================================================
def estimate_tokens(messages: List[Dict[str, Any]]) -> int:
    total_chars = 0
    for m in messages:
        total_chars += len(str(m.get("content") or ""))
        if m.get("tool_calls"):
            total_chars += len(json.dumps(m["tool_calls"]))
    return total_chars // 4

def compress_messages(messages: List[Dict[str, Any]], keep_last_rounds: int = 3) -> List[Dict[str, Any]]:
    if len(messages) < 2:
        return messages
    head = messages[:2]
    rest = messages[2:]
    tail_count = keep_last_rounds * 2
    if len(rest) <= tail_count:
        return messages

    body, tail = rest[:-tail_count], rest[-tail_count:]
    compressed_body = []
    for m in body:
        content = m.get("content") or ""
        if len(content) > 400:
            omitted = len(content) - 400
            content = content[:400] + f"\n...[{omitted} chars truncated to save context]"
            m = {**m, "content": content}
        compressed_body.append(m)
    return head + compressed_body + tail