import asyncio
import json
import logging
from datetime import datetime
from typing import AsyncGenerator, List, Dict, Any, Tuple
from pinchtab_browser import PinchTabBrowser
from web_agent_config import WEB_AGENT_SYSTEM_PROMPT, WEB_TOOL_CONTENT
import shared_utils
import app  # Still needed for generate_response and MAX_NEW_TOKENS for now

logger = logging.getLogger(__name__)

async def run_web_agent_streaming(
    goal: str,
    max_rounds: int = 15,
    inference_mode: str = "Auto (local + cloud backup)"
) -> AsyncGenerator[str, None]:
    mode = app._INFERENCE_MODE_MAP.get(inference_mode, "auto")
    browser = PinchTabBrowser()
    html_parts = [shared_utils.render_user_message(goal)]
    yield ''.join(html_parts)

    system_prompt = WEB_AGENT_SYSTEM_PROMPT.format(current_date=datetime.now().strftime('%Y-%m-%d'))
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": goal}
    ]

    tools = WEB_TOOL_CONTENT
    context_budget = await app.get_context_budget(mode)

    try:
        for round_num in range(1, max_rounds + 1):
            html_parts.append(shared_utils.render_round_badge(round_num, max_rounds))
            yield ''.join(html_parts)

            # Context management
            if shared_utils.estimate_tokens(messages) > context_budget:
                messages = shared_utils.compress_messages(messages, keep_last_rounds=3)

            html_parts.append('<div class="thinking-streaming">Processing...</div>')
            yield ''.join(html_parts)

            try:
                api_result = await app.generate_response(messages, tools, max_new_tokens=app.MAX_NEW_TOKENS, mode=mode)
                html_parts.pop() # Remove placeholder
            except Exception as e:
                html_parts.pop()
                html_parts.append(f"<p style='color:#dc2626;'>Generation Error: {str(e)}</p>")
                yield ''.join(html_parts)
                return

            generated = api_result["content"]
            reasoning, content = shared_utils.extract_thinking(generated)
            tool_call, clean_content = shared_utils.parse_tool_call(content)

            if reasoning:
                html_parts.append(shared_utils.render_thinking_collapsed(reasoning))
                yield ''.join(html_parts)

            if tool_call:
                fn_name = tool_call.get("name", "unknown")
                args = tool_call.get("arguments", {})
                html_parts.append(shared_utils.render_tool_call(fn_name, args))
                yield ''.join(html_parts)

            if clean_content.strip() and not tool_call:
                html_parts.append(f'<div class="answer-section">{clean_content}</div>')
                yield ''.join(html_parts)

            messages.append({
                "role": "assistant",
                "content": clean_content if tool_call is None else "",
                "reasoning_content": reasoning,
                "tool_calls": [{
                    "id": str(round_num),
                    "type": "function",
                    "function": {
                        "name": tool_call.get("name", ""),
                        "arguments": tool_call.get("arguments", {})
                    }
                }] if tool_call else None
            })

            if tool_call:
                fn_name = tool_call.get("name", "")
                actual_fn = fn_name.split(".", 1)[1] if "." in fn_name else fn_name
                args = tool_call.get("arguments", {})

                result = ""
                try:
                    if actual_fn == "navigate":
                        res = await browser.navigate(args.get("url"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "snapshot":
                        res = await browser.snapshot(args.get("filter", "interactive"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "click":
                        res = await browser.action("click", ref=args.get("ref"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "type":
                        res = await browser.action("type", ref=args.get("ref"), text=args.get("text"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "press":
                        res = await browser.action("press", key=args.get("key"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "get_text":
                        res = await browser.get_text()
                        result = json.dumps(res, indent=2)
                    else:
                        result = f"Unknown tool: {fn_name}"
                except Exception as e:
                    result = f"Tool error: {str(e)}"

                html_parts.append(shared_utils.render_tool_result(result, fn_name))
                yield ''.join(html_parts)

                messages.append({
                    "role": "tool",
                    "tool_call_id": str(round_num),
                    "content": result
                })
                continue

            if shared_utils.is_final_answer(generated):
                html_parts.append(shared_utils.render_completion())
                yield ''.join(html_parts)
                break
    finally:
        await browser.close()