"""Gmail RLM via ModelRelay data_source wrapper pattern."""

import sys
from dataclasses import dataclass, field
from typing import Literal, TypedDict

import httpx

from .cloud_auth import (
    CloudCredentials,
    Provider,
    WrapperSession,
    ensure_authenticated,
    get_cloud_url,
    get_wrapper_session,
)
from .logging import log_event
from .modelrelay import (
    MODELRELAY_BASE_URL,
    ModelRelayError,
    ModelResolutionPolicy,
    get_api_key,
    resolve_root_model,
)


# TypedDict definitions for data_source schema
class SchemaRequestField(TypedDict):
    """Schema for request parameters."""

    pass  # Dynamic keys


class SchemaResponseField(TypedDict):
    """Schema for response structure."""

    pass  # Dynamic keys


class SchemaEndpoint(TypedDict):
    """Schema for a single API endpoint."""

    description: str
    request: dict[str, str]
    response: dict[str, object]


class DataSourceSchema(TypedDict):
    """Schema for Gmail wrapper API."""

    search: SchemaEndpoint
    content: SchemaEndpoint


class DataSource(TypedDict):
    """ModelRelay data_source configuration."""

    type: Literal["wrapper_v1"]
    base_url: str
    token: str
    schema: DataSourceSchema


@dataclass
class EmailCitation:
    """Citation for an email referenced in the answer."""

    email_id: str
    subject: str | None = None
    sender: str | None = None
    date: str | None = None  # ISO8601


@dataclass
class GmailRLMResult:
    """Result from a Gmail RLM query."""

    answer: str
    citations: list[EmailCitation] = field(default_factory=list)
    iterations: int = 0
    total_tokens: int = 0
    error: str | None = None


def _build_gmail_schema() -> DataSourceSchema:
    """Build the wrapper_v1 schema for Gmail API.

    This schema documents the Gmail wrapper API for the LLM,
    helping it generate correct API calls.
    """
    return {
        "search": {
            "description": "Search Gmail using Gmail search syntax",
            "request": {
                "query": "string - Gmail search query (e.g. 'newer_than:7d', 'from:user@example.com', 'is:unread', 'subject:hello')",
                "page_size": "int - number of results (default 20, max 100)",
            },
            "response": {
                "items": [
                    {
                        "id": "string - email ID",
                        "subject": "string",
                        "from": {"name": "string", "email": "string"},
                        "date": "string - ISO8601 timestamp",
                        "snippet": "string - preview text",
                        "labels": ["string"],
                        "has_attachments": "boolean",
                    }
                ],
                "next_cursor": "string - for pagination",
            },
        },
        "content": {
            "description": "Get full email content by ID",
            "request": {
                "id": "string - email ID from search results",
                "format": "string - 'text' or 'html' (default 'text')",
                "max_bytes": "int - max content size (default 10000)",
            },
            "response": {
                "id": "string",
                "content": "string - email body",
                "truncated": "boolean",
            },
        },
    }


def _build_data_source(base_url: str, token: str) -> DataSource:
    """Build complete data_source config for ModelRelay.

    Args:
        base_url: Cloud backend base URL
        token: Wrapper session token

    Returns:
        DataSource config for ModelRelay RLM execute endpoint
    """
    return {
        "type": "wrapper_v1",
        "base_url": f"{base_url}/api/v1/wrapper/gmail",
        "token": token,
        "schema": _build_gmail_schema(),
    }


def run_gmail_rlm(
    question: str,
    creds: CloudCredentials,
    wrapper_session: WrapperSession | None = None,
    model: str | None = None,
    account_id: str | None = None,
    verbose: bool = False,
) -> GmailRLMResult:
    """Run an RLM query against Gmail via ModelRelay data_source wrapper.

    This calls ModelRelay's RLM execute endpoint with a data_source config
    pointing to the cloud backend's Gmail wrapper. ModelRelay handles the
    sandbox execution and automatically provides data_source_search() and
    data_source_content() functions to LLM-generated code.

    Args:
        question: The user's question about their Gmail
        creds: Cloud credentials for authentication
        wrapper_session: Pre-existing wrapper session (optional, will fetch if needed)
        model: Model to use for RLM (defaults to gemini-3-flash-preview)
        account_id: Optional specific Gmail account ID (for multi-account)
        verbose: Whether to print progress messages

    Returns:
        GmailRLMResult with answer, citations, and stats

    Raises:
        CloudAuthError: If authentication fails
        ModelRelayError: If ModelRelay API call fails
    """
    provider: Provider = "gmail"

    # Get wrapper session if not provided
    if wrapper_session is None or wrapper_session.is_expired():
        if verbose:
            print("Getting wrapper session...", file=sys.stderr)
        wrapper_session = get_wrapper_session(creds, provider=provider, account_id=account_id)

    # Resolve model
    policy = ModelResolutionPolicy(allow_env=True)
    resolved_model = resolve_root_model(model, policy=policy)

    # Build data_source config for ModelRelay
    cloud_url = get_cloud_url()
    data_source = _build_data_source(cloud_url, wrapper_session.token)

    # Build RLM request for ModelRelay
    request_body = {
        "model": resolved_model,
        "query": question,
        "data_source": data_source,
    }

    if verbose:
        print(f"Querying Gmail with model {resolved_model}...", file=sys.stderr)
        log_event(
            "gmail_rlm_request",
            model=resolved_model,
            question_length=len(question),
        )

    # Call ModelRelay RLM execute endpoint
    api_key = get_api_key()

    try:
        with httpx.Client(timeout=120.0) as client:
            response = client.post(
                f"{MODELRELAY_BASE_URL}/rlm/execute",
                headers={
                    "X-ModelRelay-Api-Key": api_key,
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                },
                json=request_body,
            )

            if response.status_code >= 400:
                error_text = response.text
                log_event(
                    "gmail_rlm_error",
                    status_code=response.status_code,
                    error=error_text,
                )
                raise ModelRelayError(response.status_code, error_text, error_text)

            data = response.json()

    except httpx.RequestError as exc:
        log_event("gmail_rlm_network_error", error=str(exc))
        raise ModelRelayError(0, f"Network error: {exc}", "")

    # Parse response - fail fast on missing required fields
    answer = data.get("answer")
    if answer is None:
        log_event("gmail_rlm_missing_answer", response_keys=list(data.keys()))
        raise ModelRelayError(0, "Missing 'answer' field in RLM response", str(data)[:500])

    iterations = data.get("iterations", 0)
    usage = data.get("usage", {})
    total_tokens = usage.get("total_tokens", 0)

    # Extract citations from trajectory if available
    citations: list[EmailCitation] = []
    # Future: parse trajectory to extract email IDs that were fetched

    if verbose:
        log_event(
            "gmail_rlm_complete",
            iterations=iterations,
            total_tokens=total_tokens,
            citations_count=len(citations),
        )

    return GmailRLMResult(
        answer=answer,
        citations=citations,
        iterations=iterations,
        total_tokens=total_tokens,
        error=None,
    )


def query_gmail(
    question: str,
    model: str | None = None,
    account_id: str | None = None,
    verbose: bool = False,
) -> GmailRLMResult:
    """High-level function to query Gmail, handling authentication automatically.

    This is the main entry point for Gmail queries from the CLI.

    Args:
        question: The user's question about their Gmail
        model: Model to use for RLM (optional)
        account_id: Optional specific Gmail account ID
        verbose: Whether to print progress messages

    Returns:
        GmailRLMResult with answer, citations, and stats

    Raises:
        CloudAuthError: If authentication fails
        ModelRelayError: If ModelRelay API call fails
    """
    # Ensure we have valid credentials (will prompt OAuth if needed)
    creds = ensure_authenticated(verbose=verbose)

    # Run the Gmail RLM query
    return run_gmail_rlm(
        question=question,
        creds=creds,
        model=model,
        account_id=account_id,
        verbose=verbose,
    )
