import asyncio
import hmac
import logging
import os

import httpx
from contextlib import asynccontextmanager
from fastapi import BackgroundTasks, FastAPI, HTTPException
from pydantic import BaseModel

import jobs
from browser_agent import run_task

logger = logging.getLogger(__name__)

MAX_CONCURRENT = 2
API_SECRET = os.getenv("API_SECRET", "")
KENSHI_CALLBACK_URL = os.getenv("KENSHI_CALLBACK_URL", "")
KENSHI_CALLBACK_SECRET = os.getenv("KENSHI_CALLBACK_SECRET", "")


@asynccontextmanager
async def lifespan(app: FastAPI):
    jobs.init_db()
    yield


app = FastAPI(lifespan=lifespan)


class TaskRequest(BaseModel):
    task: str
    chat_id: str = ""
    callback_url: str = ""
    secret: str = ""


def _verify(secret: str):
    if API_SECRET and not hmac.compare_digest(secret, API_SECRET):
        raise HTTPException(status_code=401, detail="Invalid secret")


@app.get("/")
async def root():
    return {"name": "KAZE", "status": "ok", "docs": "/docs", "health": "/health"}


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "model": "google/gemini-2.5-flash-lite",
        "jobs_running": jobs.count_running(),
    }


@app.get("/test-structured-output")
async def test_structured_output():
    """Diagnostic: test LLM with_structured_output — what browser_use 0.2.x calls internally."""
    import traceback as tb
    from pydantic import BaseModel
    from browser_agent import _make_llm, PRIMARY_MODEL

    class SimpleAnswer(BaseModel):
        answer: str

    try:
        llm = _make_llm(PRIMARY_MODEL)
        structured = llm.with_structured_output(SimpleAnswer)
        result = await structured.ainvoke("What is 2+2? Return a JSON with field 'answer'.")
        return {"status": "ok", "model": PRIMARY_MODEL, "result": result.model_dump()}
    except Exception as e:
        return {
            "status": "error",
            "model": PRIMARY_MODEL,
            "error": str(e),
            "type": type(e).__name__,
            "traceback": tb.format_exc(),
        }


@app.get("/test-llm")
async def test_llm():
    """Diagnostic: test LLM connection directly without browser_use."""
    import traceback as tb
    import importlib
    from browser_agent import _make_llm, PRIMARY_MODEL
    try:
        llm = _make_llm(PRIMARY_MODEL)
        response = await llm.ainvoke("Say 'KAZE online' in exactly those two words.")
        return {"status": "ok", "model": PRIMARY_MODEL, "response": str(response.content)}
    except Exception as e:
        logger.exception("[KAZE] /test-llm failed")
        return {"status": "error", "model": PRIMARY_MODEL, "error": str(e), "type": type(e).__name__}


@app.get("/test-browser")
async def test_browser():
    """Diagnostic: launch bare Playwright browser and navigate to example.com."""
    import traceback as tb
    try:
        from playwright.async_api import async_playwright
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=True, args=["--no-sandbox", "--disable-dev-shm-usage"])
            page = await browser.new_page()
            await page.goto("https://example.com", timeout=15000)
            title = await page.title()
            await browser.close()
        return {"status": "ok", "title": title}
    except Exception as e:
        return {"status": "error", "error": str(e), "traceback": tb.format_exc()}


@app.get("/test-agent")
async def test_agent():
    """Diagnostic: show browser_use version + run 1-step agent."""
    import traceback as tb, importlib.metadata as im
    from browser_agent import _make_llm, PRIMARY_MODEL
    info = {}
    try:
        info["browser_use_version"] = im.version("browser-use")
    except Exception:
        info["browser_use_version"] = "unknown"
    try:
        import browser_use as bu
        info["browser_use_api"] = dir(bu)
    except Exception as e:
        info["browser_use_api_error"] = str(e)

    try:
        from browser_use import Agent
        llm = _make_llm(PRIMARY_MODEL)
        # Try new 0.2.x API with explicit Browser
        try:
            from browser_use.browser.browser import Browser, BrowserConfig
            browser = Browser(config=BrowserConfig(
                headless=True,
                extra_chromium_args=["--no-sandbox", "--disable-dev-shm-usage"],
            ))
            agent = Agent(task="Navigate to example.com and return the page title.", llm=llm, browser=browser)
            info["api"] = "0.2.x (explicit Browser)"
        except ImportError:
            agent = Agent(task="Navigate to example.com and return the page title.", llm=llm)
            info["api"] = "0.1.x (auto Browser)"
        result = await agent.run(max_steps=3)
        return {"status": "ok", "result": str(result), **info}
    except Exception as e:
        return {"status": "error", "error": str(e), "traceback": tb.format_exc(), **info}


@app.post("/run")
async def run_sync(req: TaskRequest):
    _verify(req.secret)
    if jobs.count_running() >= MAX_CONCURRENT:
        raise HTTPException(status_code=503, detail="Max concurrent jobs reached. Try /run/async.")
    job_id = jobs.create_job(req.task, req.chat_id)
    jobs.update_job(job_id, "running")
    try:
        output = await run_task(req.task)
        jobs.update_job(job_id, "done", output["result"])
        return {"status": "done", "result": output["result"], "screenshots": [], "job_id": job_id}
    except asyncio.TimeoutError:
        jobs.update_job(job_id, "timeout", "Task exceeded 120s time limit")
        return {"status": "timeout", "result": "Task exceeded 120s time limit", "screenshots": [], "job_id": job_id}
    except Exception as e:
        jobs.update_job(job_id, "failed", str(e))
        return {"status": "failed", "result": str(e), "screenshots": [], "job_id": job_id}


async def _run_async_job(job_id: str, task: str, chat_id: str, callback_url: str):
    jobs.update_job(job_id, "running")
    try:
        output = await run_task(task)
        jobs.update_job(job_id, "done", output["result"])
        status, result = "done", output["result"]
    except asyncio.TimeoutError:
        jobs.update_job(job_id, "timeout", "Task exceeded 120s time limit")
        status, result = "timeout", "Task exceeded 120s time limit"
    except Exception as e:
        jobs.update_job(job_id, "failed", str(e))
        status, result = "failed", str(e)

    url = callback_url or KENSHI_CALLBACK_URL
    if url:
        payload = {
            "agentName": "KAZE",
            "taskId": job_id,
            "chatId": chat_id,
            "result": result,
            "status": status,
        }
        headers = {"X-Agent-Secret": KENSHI_CALLBACK_SECRET} if KENSHI_CALLBACK_SECRET else {}
        async with httpx.AsyncClient() as client:
            try:
                await client.post(url, json=payload, headers=headers, timeout=10)
            except Exception as e:
                logger.error("[KAZE] Webhook failed for job %s: %s", job_id, e)


@app.post("/run/async")
async def run_async_endpoint(req: TaskRequest, background_tasks: BackgroundTasks):
    _verify(req.secret)
    if jobs.count_running() >= MAX_CONCURRENT:
        raise HTTPException(status_code=503, detail="Max concurrent jobs reached.")
    job_id = jobs.create_job(req.task, req.chat_id)
    background_tasks.add_task(_run_async_job, job_id, req.task, req.chat_id, req.callback_url)
    return {"job_id": job_id, "status": "queued"}


@app.get("/job/{job_id}")
async def get_job(job_id: str):
    job = jobs.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    return job