import asyncio
import logging
import os
import re

from langchain_openai import ChatOpenAI
from browser_use import Agent

logger = logging.getLogger(__name__)
_RETRY_RE = re.compile(r'\b(429|500)\b|Failed to connect to LLM')

OPENROUTER_BASE = "https://openrouter.ai/api/v1"
PRIMARY_MODEL = "google/gemini-2.5-flash-lite"       # paid ~$0.0001/1k calls, no rate limits
FALLBACK_MODEL = "meta-llama/llama-3.3-70b-instruct:free"


def _make_llm(model: str) -> ChatOpenAI:
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        raise EnvironmentError("OPENROUTER_API_KEY environment variable is not set")
    return ChatOpenAI(model=model, api_key=api_key, base_url=OPENROUTER_BASE)


def _extract_result(history) -> str:
    if hasattr(history, "final_result"):
        r = history.final_result()
        if r:
            return str(r)
    if hasattr(history, "all_results"):
        for action_result in reversed(history.all_results):
            if getattr(action_result, "extracted_content", None):
                return str(action_result.extracted_content)
    return str(history)


async def _execute(llm: ChatOpenAI, task: str, timeout: int) -> dict:
    import time as _time
    t0 = _time.monotonic()
    print(f"[KAZE] _execute start task={task[:60]!r}", flush=True)
    agent = Agent(task=task, llm=llm)
    print("[KAZE] agent created, running (max_steps=25)...", flush=True)
    try:
        history = await asyncio.wait_for(agent.run(max_steps=25), timeout=timeout)
    except Exception:
        print(f"[KAZE] agent.run() raised after {_time.monotonic()-t0:.1f}s", flush=True)
        logger.exception("[KAZE] agent.run() raised")
        raise
    n_r = len(history.all_results) if hasattr(history, "all_results") else "?"
    n_o = len(history.all_model_outputs) if hasattr(history, "all_model_outputs") else "?"
    final = history.final_result() if hasattr(history, "final_result") else None
    elapsed = _time.monotonic() - t0
    print(f"[KAZE] run done in {elapsed:.1f}s: results={n_r} outputs={n_o} final={final!r}", flush=True)
    logger.info("[KAZE] run done: results=%s outputs=%s final=%r", n_r, n_o, final)
    return {"result": _extract_result(history), "screenshots": []}


async def run_task(task: str, timeout: int = 120) -> dict:
    llm = _make_llm(PRIMARY_MODEL)
    try:
        return await _execute(llm, task, timeout)
    except Exception as e:
        logger.exception("[KAZE] primary model failed: %s", e)
        if _RETRY_RE.search(str(e)):
            logger.info("[KAZE] retrying with fallback model")
            llm = _make_llm(FALLBACK_MODEL)
            return await _execute(llm, task, timeout)
        raise