import json
import logging
import os
import platform
import queue
import re
import threading
import time
import uuid
from collections.abc import Iterator
from contextlib import asynccontextmanager
from dataclasses import asdict
from typing import Any

import gradio as gr
import torch
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from transformers import (
    AutoModelForMultimodalLM,
    AutoProcessor,
    TextIteratorStreamer,
)

from observability import (
    GenerationStats,
    RuntimeMetrics,
    configure_logger,
    env_flag,
    log_event,
    summarize_messages,
)
from prompting import (
    build_gemma4_prompt,
    normalize_openai_messages,
    parse_gemma4_arguments,
    prepare_messages,
)
from space_ui import (
    CHAT_PLACEHOLDER,
    EXAMPLE_LABELS,
    EXAMPLES,
    SPACE_CSS,
    SPACE_HEAD,
    space_description,
)

MODEL_ID = os.getenv("MODEL_ID", "google/gemma-4-E2B-it")
HF_TOKEN = os.getenv("HF_TOKEN")
MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "8192"))
DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", "512"))
LOG_PAYLOADS = env_flag("LOG_PAYLOADS")
LOG_STREAM_CHUNKS = env_flag("LOG_STREAM_CHUNKS")
LOG_MAX_TEXT_CHARS = int(os.getenv("LOG_MAX_TEXT_CHARS", "2000"))
SERVER_STARTED_AT = time.time()
logger = configure_logger()

generation_lock = threading.Lock()

log_event(
    logger,
    "server.boot",
    python=platform.python_version(),
    torch=torch.__version__,
    cuda_available=torch.cuda.is_available(),
    model=MODEL_ID,
    max_input_tokens=MAX_INPUT_TOKENS,
    default_max_tokens=DEFAULT_MAX_TOKENS,
    log_payloads=LOG_PAYLOADS,
    log_stream_chunks=LOG_STREAM_CHUNKS,
)

processor_started = time.perf_counter()
log_event(logger, "model.processor.load.started", model=MODEL_ID)
try:
    processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
except Exception:
    log_event(
        logger,
        "model.processor.load.failed",
        level=logging.ERROR,
        exc_info=True,
        model=MODEL_ID,
    )
    raise
log_event(
    logger,
    "model.processor.load.completed",
    model=MODEL_ID,
    duration_ms=round((time.perf_counter() - processor_started) * 1000, 2),
    has_chat_template=bool(processor.chat_template),
)

prompt_mode = "processor-template" if processor.chat_template else "gemma4-fallback"
model_started = time.perf_counter()
log_event(
    logger,
    "model.load.started",
    model=MODEL_ID,
    device="cuda" if torch.cuda.is_available() else "cpu",
)
try:
    if torch.cuda.is_available():
        model = AutoModelForMultimodalLM.from_pretrained(
            MODEL_ID,
            token=HF_TOKEN,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            low_cpu_mem_usage=True,
        )
    else:
        model = AutoModelForMultimodalLM.from_pretrained(
            MODEL_ID,
            token=HF_TOKEN,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
        )
except Exception:
    log_event(
        logger,
        "model.load.failed",
        level=logging.ERROR,
        exc_info=True,
        model=MODEL_ID,
    )
    raise

model.eval()
log_event(
    logger,
    "model.load.completed",
    model=MODEL_ID,
    duration_ms=round((time.perf_counter() - model_started) * 1000, 2),
    device=str(model.device),
    prompt_mode=prompt_mode,
)


metrics = RuntimeMetrics(SERVER_STARTED_AT)


class ChatCompletionRequest(BaseModel):
    model: str = MODEL_ID
    messages: list[dict[str, Any]]
    tools: list[dict[str, Any]] | None = None
    tool_choice: Any | None = None
    stream: bool = False
    temperature: float = Field(default=0.7, ge=0)
    top_p: float = Field(default=0.9, gt=0, le=1)
    max_tokens: int = Field(default=DEFAULT_MAX_TOKENS, ge=1, le=4096)


def build_inputs(
    request: ChatCompletionRequest,
    request_id: str,
) -> dict[str, torch.Tensor]:
    messages = prepare_messages(normalize_openai_messages(request.messages))
    build_started = time.perf_counter()
    if processor.chat_template:
        template_kwargs: dict[str, Any] = {
            "conversation": messages,
            "add_generation_prompt": True,
            "enable_thinking": False,
            "tokenize": True,
            "return_dict": True,
            "return_tensors": "pt",
            "truncation": True,
            "max_length": MAX_INPUT_TOKENS,
        }
        if request.tools:
            template_kwargs["tools"] = request.tools
        inputs = processor.apply_chat_template(**template_kwargs)
    else:
        prompt = build_gemma4_prompt(
            messages,
            tools=request.tools,
            bos_token=processor.tokenizer.bos_token or "",
        )
        inputs = processor(
            text=prompt,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_INPUT_TOKENS,
        )

    device_inputs = {key: value.to(model.device) for key, value in inputs.items()}
    prompt_tokens = int(device_inputs["input_ids"].shape[-1])
    log_event(
        logger,
        "prompt.built",
        request_id=request_id,
        prompt_mode=prompt_mode,
        prompt_tokens=prompt_tokens,
        input_truncated=prompt_tokens >= MAX_INPUT_TOKENS,
        tool_count=len(request.tools or []),
        duration_ms=round((time.perf_counter() - build_started) * 1000, 2),
        **summarize_messages(messages),
    )
    if LOG_PAYLOADS:
        log_event(
            logger,
            "prompt.payload",
            request_id=request_id,
            messages=messages,
            tools=request.tools,
            max_text_chars=LOG_MAX_TEXT_CHARS,
        )
    return device_inputs


def generation_kwargs(
    request: ChatCompletionRequest,
    inputs: dict[str, torch.Tensor],
) -> dict[str, Any]:
    kwargs: dict[str, Any] = {
        **inputs,
        "max_new_tokens": request.max_tokens,
        "do_sample": request.temperature > 0,
        "pad_token_id": processor.tokenizer.pad_token_id,
    }
    if request.temperature > 0:
        kwargs["temperature"] = request.temperature
        kwargs["top_p"] = request.top_p
    return kwargs


def count_tokens(text: str) -> int:
    if not text:
        return 0
    return len(
        processor.tokenizer.encode(
            text,
            add_special_tokens=False,
        )
    )


def parse_tool_call(text: str) -> tuple[str, dict[str, Any]] | None:
    candidates = [
        match.group(1)
        for match in re.finditer(
            r"<\|tool_call>\s*(.*?)\s*(?:<tool_call\|>|<turn\|>|$)",
            text,
            re.DOTALL,
        )
    ]
    candidates.append(text.strip())

    for candidate in candidates:
        candidate = (
            candidate.strip()
            .removeprefix("```json")
            .removesuffix("```")
            .replace('<|"|>', '"')
            .strip()
        )
        native_call = re.fullmatch(
            r"call:(?P<name>[A-Za-z_][A-Za-z0-9_]*)(?P<arguments>\{.*\})",
            candidate,
            re.DOTALL,
        )
        if native_call:
            arguments = parse_gemma4_arguments(native_call.group("arguments"))
            if arguments is not None:
                return native_call.group("name"), arguments
            continue

        try:
            parsed = json.loads(candidate)
        except json.JSONDecodeError:
            continue

        if isinstance(parsed, dict):
            name = parsed.get("name") or parsed.get("function")
            arguments = parsed.get("arguments") or parsed.get("parameters") or {}
            if isinstance(name, str) and isinstance(arguments, dict):
                return name, arguments
    return None


def generate(
    request: ChatCompletionRequest,
    request_id: str,
) -> tuple[str, GenerationStats]:
    stats = GenerationStats()
    started = time.perf_counter()
    metrics.generation_started()
    failed = False

    try:
        inputs = build_inputs(request, request_id)
        input_length = int(inputs["input_ids"].shape[-1])
        stats.prompt_tokens = input_length
        queued_at = time.perf_counter()
        with generation_lock:
            stats.queue_ms = round((time.perf_counter() - queued_at) * 1000, 2)
            inference_started = time.perf_counter()
            with torch.inference_mode():
                output = model.generate(**generation_kwargs(request, inputs))
            stats.inference_ms = round(
                (time.perf_counter() - inference_started) * 1000,
                2,
            )

        text = processor.tokenizer.decode(
            output[0][input_length:],
            skip_special_tokens=not bool(request.tools),
        ).strip()
        stats.completion_tokens = int(output[0].shape[-1]) - input_length
        stats.output_chars = len(text)
        stats.chunks = 1
        log_event(
            logger,
            "generation.completed",
            request_id=request_id,
            stream=False,
            total_ms=round((time.perf_counter() - started) * 1000, 2),
            **asdict(stats),
        )
        if LOG_PAYLOADS:
            log_event(
                logger,
                "generation.output",
                request_id=request_id,
                text=text,
                max_text_chars=LOG_MAX_TEXT_CHARS,
            )
        return text, stats
    except Exception:
        failed = True
        log_event(
            logger,
            "generation.failed",
            level=logging.ERROR,
            exc_info=True,
            request_id=request_id,
            stream=False,
            total_ms=round((time.perf_counter() - started) * 1000, 2),
            **asdict(stats),
        )
        raise
    finally:
        metrics.generation_finished(stats, failed=failed)


def stream_generate(
    request: ChatCompletionRequest,
    request_id: str,
    stats: GenerationStats,
) -> Iterator[str]:
    started = time.perf_counter()
    output_parts: list[str] = []
    failed = False
    metrics.generation_started()

    try:
        inputs = build_inputs(request, request_id)
        stats.prompt_tokens = int(inputs["input_ids"].shape[-1])
        streamer = TextIteratorStreamer(
            processor.tokenizer,
            skip_prompt=True,
            skip_special_tokens=not bool(request.tools),
            timeout=120,
        )
        kwargs = generation_kwargs(request, inputs)
        kwargs["streamer"] = streamer
        thread_errors: list[BaseException] = []

        def run_model() -> None:
            try:
                with torch.inference_mode():
                    model.generate(**kwargs)
            except BaseException as error:
                thread_errors.append(error)

        thread = threading.Thread(target=run_model, daemon=True)
        queued_at = time.perf_counter()
        with generation_lock:
            stats.queue_ms = round((time.perf_counter() - queued_at) * 1000, 2)
            inference_started = time.perf_counter()
            thread.start()
            try:
                for text in streamer:
                    if text and stats.first_token_ms is None:
                        stats.first_token_ms = round(
                            (time.perf_counter() - inference_started) * 1000,
                            2,
                        )
                    stats.chunks += 1
                    stats.output_chars += len(text)
                    output_parts.append(text)
                    if LOG_STREAM_CHUNKS:
                        log_event(
                            logger,
                            "generation.stream.chunk",
                            request_id=request_id,
                            chunk_index=stats.chunks,
                            chars=len(text),
                            text=text if LOG_PAYLOADS else None,
                            max_text_chars=LOG_MAX_TEXT_CHARS,
                        )
                    yield text
            except queue.Empty as error:
                if thread_errors:
                    raise RuntimeError("Model generation failed") from thread_errors[0]
                raise TimeoutError("Timed out waiting for the next model token") from error
            finally:
                thread.join()
                stats.inference_ms = round(
                    (time.perf_counter() - inference_started) * 1000,
                    2,
                )

        if thread_errors:
            raise RuntimeError("Model generation failed") from thread_errors[0]

        output_text = "".join(output_parts)
        stats.completion_tokens = count_tokens(output_text)
        log_event(
            logger,
            "generation.completed",
            request_id=request_id,
            stream=True,
            total_ms=round((time.perf_counter() - started) * 1000, 2),
            **asdict(stats),
        )
        if LOG_PAYLOADS:
            log_event(
                logger,
                "generation.output",
                request_id=request_id,
                text=output_text,
                max_text_chars=LOG_MAX_TEXT_CHARS,
            )
    except Exception:
        failed = True
        log_event(
            logger,
            "generation.failed",
            level=logging.ERROR,
            exc_info=True,
            request_id=request_id,
            stream=True,
            total_ms=round((time.perf_counter() - started) * 1000, 2),
            **asdict(stats),
        )
        raise
    finally:
        if not stats.completion_tokens and output_parts:
            stats.completion_tokens = count_tokens("".join(output_parts))
        metrics.generation_finished(stats, failed=failed)


def chunk(
    completion_id: str,
    model_name: str,
    delta: dict[str, Any],
    finish_reason: str | None = None,
) -> str:
    payload = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model_name,
        "choices": [
            {
                "index": 0,
                "delta": delta,
                "finish_reason": finish_reason,
            }
        ],
    }
    return f"data: {json.dumps(payload)}\n\n"


@asynccontextmanager
async def lifespan(_: FastAPI):
    log_event(
        logger,
        "server.ready",
        model=MODEL_ID,
        device=str(model.device),
        prompt_mode=prompt_mode,
        docs="/docs",
        chat_completions="/v1/chat/completions",
    )
    yield
    log_event(logger, "server.shutdown", **metrics.snapshot())


api = FastAPI(
    title="Browso Agent",
    description=(
        "OpenAI-compatible Gemma 4 inference server for browser assistance "
        "and structured tool use."
    ),
    version="0.2.0",
    lifespan=lifespan,
)


@api.middleware("http")
async def log_http_request(http_request: Request, call_next):
    request_id = http_request.headers.get("x-request-id") or f"req-{uuid.uuid4().hex}"
    http_request.state.request_id = request_id
    started = time.perf_counter()
    client_host = http_request.client.host if http_request.client else None
    log_event(
        logger,
        "http.request.started",
        request_id=request_id,
        method=http_request.method,
        path=http_request.url.path,
        query=http_request.url.query or None,
        client=client_host,
        user_agent=http_request.headers.get("user-agent"),
        content_length=http_request.headers.get("content-length"),
    )

    try:
        response = await call_next(http_request)
    except Exception:
        metrics.record_http(500)
        log_event(
            logger,
            "http.request.failed",
            level=logging.ERROR,
            exc_info=True,
            request_id=request_id,
            method=http_request.method,
            path=http_request.url.path,
            duration_ms=round((time.perf_counter() - started) * 1000, 2),
        )
        raise

    response.headers["x-request-id"] = request_id
    metrics.record_http(response.status_code)
    log_event(
        logger,
        "http.request.completed",
        request_id=request_id,
        method=http_request.method,
        path=http_request.url.path,
        status_code=response.status_code,
        duration_ms=round((time.perf_counter() - started) * 1000, 2),
    )
    return response


@api.get("/health")
def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "model": MODEL_ID,
        "modelReady": True,
        "promptMode": prompt_mode,
        "device": str(model.device),
        "uptimeSeconds": metrics.snapshot()["uptime_seconds"],
    }


@api.get("/status")
def status() -> dict[str, Any]:
    return {
        "status": "ok",
        "model": MODEL_ID,
        "device": str(model.device),
        "promptMode": prompt_mode,
        "logging": {
            "level": logging.getLevelName(logger.level),
            "payloads": LOG_PAYLOADS,
            "streamChunks": LOG_STREAM_CHUNKS,
        },
        "metrics": metrics.snapshot(),
    }


@api.get("/v1/models")
def models() -> dict[str, Any]:
    return {
        "object": "list",
        "data": [{"id": MODEL_ID, "object": "model", "owned_by": "browso"}],
    }


@api.post("/v1/chat/completions")
def chat_completions(
    payload: ChatCompletionRequest,
    http_request: Request,
):
    request_id = http_request.state.request_id
    if not payload.messages:
        log_event(
            logger,
            "completion.rejected",
            level=logging.WARNING,
            request_id=request_id,
            reason="messages must not be empty",
        )
        raise HTTPException(status_code=400, detail="messages must not be empty")

    completion_id = f"chatcmpl-{uuid.uuid4().hex}"
    log_event(
        logger,
        "completion.requested",
        request_id=request_id,
        completion_id=completion_id,
        requested_model=payload.model,
        backend_model=MODEL_ID,
        stream=payload.stream,
        temperature=payload.temperature,
        top_p=payload.top_p,
        max_tokens=payload.max_tokens,
        tool_count=len(payload.tools or []),
        tool_choice=payload.tool_choice,
        **summarize_messages(payload.messages),
    )
    if LOG_PAYLOADS:
        log_event(
            logger,
            "completion.payload",
            request_id=request_id,
            completion_id=completion_id,
            payload=payload.model_dump(),
            max_text_chars=LOG_MAX_TEXT_CHARS,
        )

    if payload.stream:
        def event_stream() -> Iterator[str]:
            stats = GenerationStats()
            finish_reason = "stop"
            try:
                yield chunk(completion_id, payload.model, {"role": "assistant"})
                if payload.tools:
                    raw_text = "".join(
                        stream_generate(payload, request_id, stats)
                    )
                    tool_call = parse_tool_call(raw_text)
                    if tool_call:
                        name, arguments = tool_call
                        finish_reason = "tool_calls"
                        log_event(
                            logger,
                            "tool_call.parsed",
                            request_id=request_id,
                            completion_id=completion_id,
                            name=name,
                            argument_keys=sorted(arguments),
                        )
                        if LOG_PAYLOADS:
                            log_event(
                                logger,
                                "tool_call.payload",
                                request_id=request_id,
                                completion_id=completion_id,
                                name=name,
                                arguments=arguments,
                                max_text_chars=LOG_MAX_TEXT_CHARS,
                            )
                        yield chunk(
                            completion_id,
                            payload.model,
                            {
                                "tool_calls": [
                                    {
                                        "index": 0,
                                        "id": f"call_{uuid.uuid4().hex}",
                                        "type": "function",
                                        "function": {
                                            "name": name,
                                            "arguments": json.dumps(arguments),
                                        },
                                    }
                                ]
                            },
                        )
                        yield chunk(
                            completion_id,
                            payload.model,
                            {},
                            finish_reason,
                        )
                    else:
                        log_event(
                            logger,
                            "tool_call.not_detected",
                            level=logging.WARNING,
                            request_id=request_id,
                            completion_id=completion_id,
                            output_chars=len(raw_text),
                        )
                        yield chunk(
                            completion_id,
                            payload.model,
                            {"content": raw_text},
                        )
                        yield chunk(completion_id, payload.model, {}, "stop")
                else:
                    for text in stream_generate(payload, request_id, stats):
                        yield chunk(
                            completion_id,
                            payload.model,
                            {"content": text},
                        )
                    yield chunk(completion_id, payload.model, {}, "stop")
                yield "data: [DONE]\n\n"
                log_event(
                    logger,
                    "completion.stream.completed",
                    request_id=request_id,
                    completion_id=completion_id,
                    finish_reason=finish_reason,
                    **asdict(stats),
                )
            except GeneratorExit:
                log_event(
                    logger,
                    "completion.stream.disconnected",
                    level=logging.WARNING,
                    request_id=request_id,
                    completion_id=completion_id,
                    **asdict(stats),
                )
                raise
            except Exception:
                log_event(
                    logger,
                    "completion.stream.failed",
                    level=logging.ERROR,
                    exc_info=True,
                    request_id=request_id,
                    completion_id=completion_id,
                    **asdict(stats),
                )
                raise

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    try:
        raw_text, stats = generate(payload, request_id)
    except Exception as error:
        raise HTTPException(
            status_code=500,
            detail=f"Inference failed. Reference request ID {request_id}.",
        ) from error
    tool_call = parse_tool_call(raw_text) if payload.tools else None
    message: dict[str, Any] = {"role": "assistant", "content": raw_text}
    finish_reason = "stop"
    if tool_call:
        name, arguments = tool_call
        log_event(
            logger,
            "tool_call.parsed",
            request_id=request_id,
            completion_id=completion_id,
            name=name,
            argument_keys=sorted(arguments),
        )
        message = {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": f"call_{uuid.uuid4().hex}",
                    "type": "function",
                    "function": {
                        "name": name,
                        "arguments": json.dumps(arguments),
                    },
                }
            ],
        }
        finish_reason = "tool_calls"

    log_event(
        logger,
        "completion.completed",
        request_id=request_id,
        completion_id=completion_id,
        finish_reason=finish_reason,
        **asdict(stats),
    )
    return {
        "id": completion_id,
        "object": "chat.completion",
        "created": int(time.time()),
        "model": payload.model,
        "choices": [
            {
                "index": 0,
                "message": message,
                "finish_reason": finish_reason,
            }
        ],
        "usage": {
            "prompt_tokens": stats.prompt_tokens,
            "completion_tokens": stats.completion_tokens,
            "total_tokens": stats.prompt_tokens + stats.completion_tokens,
        },
    }


def respond(message: str, history: list[dict[str, Any]]) -> Iterator[str]:
    request_id = f"ui-{uuid.uuid4().hex}"
    messages = list(history or [])
    messages.append({"role": "user", "content": message})
    payload = ChatCompletionRequest(messages=messages, stream=True)
    stats = GenerationStats()
    response = ""
    log_event(
        logger,
        "ui.chat.requested",
        request_id=request_id,
        **summarize_messages(messages),
    )
    if LOG_PAYLOADS:
        log_event(
            logger,
            "ui.chat.payload",
            request_id=request_id,
            messages=messages,
            max_text_chars=LOG_MAX_TEXT_CHARS,
        )

    try:
        for text in stream_generate(payload, request_id, stats):
            response += text
            yield response
        log_event(
            logger,
            "ui.chat.completed",
            request_id=request_id,
            **asdict(stats),
        )
    except GeneratorExit:
        log_event(
            logger,
            "ui.chat.disconnected",
            level=logging.WARNING,
            request_id=request_id,
            **asdict(stats),
        )
        raise
    except Exception:
        log_event(
            logger,
            "ui.chat.failed",
            level=logging.ERROR,
            exc_info=True,
            request_id=request_id,
            **asdict(stats),
        )
        raise


demo = gr.ChatInterface(
    fn=respond,
    chatbot=gr.Chatbot(
        label="Browso Agent",
        show_label=False,
        height="58vh",
        min_height=420,
        layout="panel",
        buttons=["copy", "copy_all"],
        watermark="Generated by Browso Agent",
        placeholder=CHAT_PLACEHOLDER,
        elem_classes=["browso-chat"],
    ),
    textbox=gr.Textbox(
        placeholder="Ask Browso Agent about a page, a task, code, or research...",
        show_label=False,
        container=False,
        lines=1,
        max_lines=8,
        autofocus=True,
        submit_btn="Send",
        stop_btn="Stop",
        elem_classes=["browso-input"],
    ),
    description=space_description(MODEL_ID, prompt_mode),
    examples=EXAMPLES,
    example_labels=EXAMPLE_LABELS,
    run_examples_on_click=True,
    cache_examples=False,
    flagging_mode="never",
    fill_height=False,
    fill_width=True,
    api_name="chat",
    api_description="Stream a response from the Browso Agent model.",
    save_history=True,
)
demo.title = "Browso Agent"

app = gr.mount_gradio_app(
    api,
    demo,
    path="/",
    footer_links=["api", "settings"],
    show_error=True,
    enable_monitoring=True,
    theme="soft",
    css=SPACE_CSS,
    head=SPACE_HEAD,
)


if __name__ == "__main__":
    log_event(
        logger,
        "uvicorn.starting",
        host="0.0.0.0",
        port=7860,
        log_level=os.getenv("UVICORN_LOG_LEVEL", "info"),
    )
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,
        log_level=os.getenv("UVICORN_LOG_LEVEL", "info"),
        access_log=False,
    )