"""Temporary high-level library of the Google GenerativeAI API.

(The content of this file should eventually go into the Python package
`google.generativeai`)
"""

import datetime
import logging
import re
from collections.abc import MutableSequence
from dataclasses import dataclass
from typing import Any
from urllib.parse import urlparse

import google.ai.generativelanguage as genai
import langchain_core
from google.ai.generativelanguage_v1beta import (
    GenerativeServiceAsyncClient as v1betaGenerativeServiceAsyncClient,
)
from google.ai.generativelanguage_v1beta import (
    GenerativeServiceClient as v1betaGenerativeServiceClient,
)
from google.api_core import client_options as client_options_lib
from google.api_core import exceptions as gapi_exception
from google.api_core import gapic_v1
from google.auth import credentials, exceptions
from google.protobuf import timestamp_pb2

_logger = logging.getLogger(__name__)
_DEFAULT_API_ENDPOINT = "generativelanguage.googleapis.com"
_USER_AGENT = f"langchain/{langchain_core.__version__}"
_DEFAULT_PAGE_SIZE = 20
_MAX_REQUEST_PER_CHUNK = 100
_NAME_REGEX = re.compile(r"^corpora/([^/]+?)(/documents/([^/]+?)(/chunks/([^/]+?))?)?$")


@dataclass
class EntityName:
    corpus_id: str
    document_id: str | None = None
    chunk_id: str | None = None

    def __post_init__(self) -> None:
        if self.chunk_id is not None and self.document_id is None:
            msg = f"Chunk must have document ID but found {self}"
            raise ValueError(msg)

    @classmethod
    def from_str(cls, encoded: str) -> "EntityName":
        matched = _NAME_REGEX.match(encoded)
        if not matched:
            msg = f"Invalid entity name: {encoded}"
            raise ValueError(msg)

        return cls(
            corpus_id=matched.group(1),
            document_id=matched.group(3),
            chunk_id=matched.group(5),
        )

    def __repr__(self) -> str:
        name = f"corpora/{self.corpus_id}"
        if self.document_id is None:
            return name
        name += f"/documents/{self.document_id}"
        if self.chunk_id is None:
            return name
        name += f"/chunks/{self.chunk_id}"
        return name

    def __str__(self) -> str:
        return repr(self)

    def is_corpus(self) -> bool:
        return self.document_id is None

    def is_document(self) -> bool:
        return self.document_id is not None and self.chunk_id is None

    def is_chunk(self) -> bool:
        return self.chunk_id is not None


@dataclass
class Corpus:
    name: str
    display_name: str | None
    create_time: timestamp_pb2.Timestamp | None
    update_time: timestamp_pb2.Timestamp | None

    @property
    def corpus_id(self) -> str:
        name = EntityName.from_str(self.name)
        return name.corpus_id

    @classmethod
    def from_corpus(cls, c: genai.Corpus) -> "Corpus":
        return cls(
            name=c.name,
            display_name=c.display_name,
            create_time=c.create_time,
            update_time=c.update_time,
        )


@dataclass
class Document:
    name: str
    display_name: str | None
    create_time: timestamp_pb2.Timestamp | None
    update_time: timestamp_pb2.Timestamp | None
    custom_metadata: MutableSequence[genai.CustomMetadata] | None

    @property
    def corpus_id(self) -> str:
        name = EntityName.from_str(self.name)
        return name.corpus_id

    @property
    def document_id(self) -> str:
        name = EntityName.from_str(self.name)
        assert isinstance(name.document_id, str)
        return name.document_id

    @classmethod
    def from_document(cls, d: genai.Document) -> "Document":
        return cls(
            name=d.name,
            display_name=d.display_name,
            create_time=d.create_time,
            update_time=d.update_time,
            custom_metadata=d.custom_metadata,
        )


@dataclass
class Config:
    """Global configuration for Google Generative AI API.

    Normally, the defaults should work fine. Use this to pass Google Auth credentials
    such as using a service account.

    Refer to for [auth credentials documentation](https://developers.google.com/identity/protocols/oauth2/service-account#creatinganaccount)

    Attributes:
        api_endpoint: The Google Generative API endpoint address.
        user_agent: The user agent to use for logging.
        page_size: For paging RPCs, how many entities to return per RPC.
        testing: Are the unit tests running?
        auth_credentials: For setting credentials such as using service accounts.
    """

    api_endpoint: str = _DEFAULT_API_ENDPOINT
    user_agent: str = _USER_AGENT
    page_size: int = _DEFAULT_PAGE_SIZE
    testing: bool = False
    auth_credentials: credentials.Credentials | None = None


def set_config(config: Config) -> None:
    """Set global defaults for operations with Google Generative AI API."""
    global _config
    _config = config


def get_config() -> Config:
    return _config


_config = Config()


class TestCredentials(credentials.Credentials):
    """Credentials that do not provide any authentication information.

    Useful for unit tests where the credentials are not used.
    """

    @property
    def expired(self) -> bool:
        """Returns `False`, test credentials never expire."""
        return False

    @property
    def valid(self) -> bool:
        """Returns `True`, test credentials are always valid."""
        return True

    def refresh(self, request: Any) -> None:
        """Raises `InvalidOperation`, test credentials cannot be refreshed."""
        msg = "Test credentials cannot be refreshed."
        # TODO: remove ignore once google-auth has types.
        raise exceptions.InvalidOperation(msg)  # type: ignore[no-untyped-call]

    def apply(self, headers: Any, token: Any = None) -> None:
        """Anonymous credentials do nothing to the request.

        The optional `token` argument is not supported.

        Raises:
            google.auth.exceptions.InvalidValue: If a token was specified.
        """
        if token is not None:
            msg = "Test credentials don't support tokens."
            # TODO: remove ignore once google-auth has types.
            raise exceptions.InvalidValue(msg)  # type: ignore[no-untyped-call]

    def before_request(self, request: Any, method: Any, url: Any, headers: Any) -> None:
        """Test credentials do nothing to the request."""


def _get_credentials() -> credentials.Credentials | None:
    """Returns credential from config if set or fake credentials for unit testing.

    If `_config.testing` is `True`, a fake credential is returned.

    Otherwise, we are in a real environment and will use credentials if provided or
    `None` is returned.

    If `None` is passed to the clients later on, the actual credentials will be
    inferred by the rules specified in google.auth package.
    """
    if _config.testing:
        # TODO: remove ignore once google-auth has types.
        return TestCredentials()  # type: ignore[no-untyped-call]
    if _config.auth_credentials:
        return _config.auth_credentials
    return None


def build_semantic_retriever() -> genai.RetrieverServiceClient:
    """Uses the default `'grpc'` transport to build a semantic retriever client."""
    credentials = _get_credentials()
    return genai.RetrieverServiceClient(
        credentials=credentials,
        # TODO: remove ignore once google-auth has types.
        client_info=gapic_v1.client_info.ClientInfo(user_agent=_USER_AGENT),  # type: ignore[no-untyped-call]
        client_options=client_options_lib.ClientOptions(
            api_endpoint=_config.api_endpoint
        ),
    )


def _prepare_config(
    credentials: credentials.Credentials | None = None,
    api_key: str | None = None,
    client_options: dict[str, Any] | None = None,
    client_info: gapic_v1.client_info.ClientInfo | None = None,
    transport: str | None = None,
) -> dict[str, Any]:
    formatted_client_options: dict = {"api_endpoint": _config.api_endpoint}
    if client_options:
        formatted_client_options.update(**client_options)

    # Validate and format API endpoint for gRPC transports
    if (
        transport in ("grpc", "grpc_asyncio")
        and "api_endpoint" in formatted_client_options
    ):
        api_endpoint = formatted_client_options["api_endpoint"]
        if isinstance(api_endpoint, str):
            _validate_grpc_endpoint(api_endpoint, transport)
            # Format the endpoint for gRPC if needed
            formatted_client_options["api_endpoint"] = _format_grpc_endpoint(
                api_endpoint
            )

    if not credentials and api_key:
        formatted_client_options["api_key"] = api_key
    elif not credentials and not api_key:
        credentials = _get_credentials()
    client_info = (
        client_info
        if client_info
        # TODO: remove ignore once google-auth has types.
        else gapic_v1.client_info.ClientInfo(user_agent=_USER_AGENT)  # type: ignore[no-untyped-call]
    )
    config = {
        "credentials": credentials,
        "client_info": client_info,
        "client_options": client_options_lib.ClientOptions(**formatted_client_options),
        "transport": transport,
    }
    return {k: v for k, v in config.items() if v is not None}


def _validate_grpc_endpoint(endpoint: str, transport: str) -> None:
    """Validate that an endpoint is compatible with gRPC transports.

    Args:
        endpoint: The API endpoint to validate.
        transport: The transport type (`'grpc'` or `'grpc_asyncio'`).

    Raises:
        ValueError: If the endpoint format is incompatible with gRPC.
    """
    if endpoint.startswith(("http://", "https://")):
        parsed = urlparse(endpoint)
        if parsed.path and parsed.path.rstrip("/") != "":
            msg = (
                f"gRPC transport '{transport}' does not support URL paths. "
                f"Endpoint '{endpoint}' contains path '{parsed.path}'. "
                "Use format 'hostname:port' or consider using 'rest' transport "
                "instead."
            )
            raise ValueError(msg)
    elif "/" in endpoint and not endpoint.startswith(
        "["
    ):  # IPv6 addresses use brackets
        msg = (
            f"gRPC transport '{transport}' does not support URL paths. "
            f"Endpoint '{endpoint}' appears to contain a path. "
            f"Use format 'hostname:port' or consider using 'rest' transport instead."
        )
        raise ValueError(msg)


def _format_grpc_endpoint(endpoint: str) -> str:
    """Format an endpoint for gRPC compatibility.

    Args:
        endpoint: The original endpoint.

    Returns:
        Properly formatted endpoint for gRPC (`hostname:port` format).
    """
    if not endpoint.startswith(("http://", "https://")) and ":" in endpoint:
        return endpoint

    # Handle URLs with protocol scheme (but no path, validated above)
    if endpoint.startswith(("http://", "https://")):
        parsed = urlparse(endpoint)
        hostname = parsed.hostname
        port = parsed.port
        if not port:
            port = 443 if endpoint.startswith("https://") else 80
        return f"{hostname}:{port}"

    # If no port specified, add default HTTPS port
    if ":" not in endpoint:
        return f"{endpoint}:443"

    return endpoint


def build_generative_service(
    credentials: credentials.Credentials | None = None,
    api_key: str | None = None,
    client_options: dict[str, Any] | None = None,
    client_info: gapic_v1.client_info.ClientInfo | None = None,
    transport: str | None = None,
) -> v1betaGenerativeServiceClient:
    config = _prepare_config(
        credentials=credentials,
        api_key=api_key,
        client_options=client_options,
        transport=transport,
        client_info=client_info,
    )
    return v1betaGenerativeServiceClient(**config)


def build_generative_async_service(
    credentials: credentials.Credentials | None,
    api_key: str | None = None,
    client_options: dict[str, Any] | None = None,
    client_info: gapic_v1.client_info.ClientInfo | None = None,
    transport: str | None = None,
) -> v1betaGenerativeServiceAsyncClient:
    config = _prepare_config(
        credentials=credentials,
        api_key=api_key,
        client_options=client_options,
        transport=transport,
        client_info=client_info,
    )
    return v1betaGenerativeServiceAsyncClient(**config)


def get_corpus(
    *,
    corpus_id: str,
    client: genai.RetrieverServiceClient,
) -> Corpus | None:
    try:
        corpus = client.get_corpus(
            genai.GetCorpusRequest(name=str(EntityName(corpus_id=corpus_id)))
        )
        return Corpus.from_corpus(corpus)
    except Exception as e:
        # If the corpus does not exist, the server returns a permission error.
        if not isinstance(e, gapi_exception.PermissionDenied):
            raise
        _logger.warning(f"Corpus {corpus_id} not found: {e}")
        return None


def create_corpus(
    *,
    corpus_id: str | None = None,
    display_name: str | None = None,
    client: genai.RetrieverServiceClient,
) -> Corpus:
    name: str | None
    name = str(EntityName(corpus_id=corpus_id)) if corpus_id is not None else None

    new_display_name = display_name or f"Untitled {datetime.datetime.now()}"

    new_corpus = client.create_corpus(
        genai.CreateCorpusRequest(
            corpus=genai.Corpus(name=name, display_name=new_display_name)
        )
    )

    return Corpus.from_corpus(new_corpus)


def delete_corpus(
    *,
    corpus_id: str,
    client: genai.RetrieverServiceClient,
) -> None:
    client.delete_corpus(
        genai.DeleteCorpusRequest(name=str(EntityName(corpus_id=corpus_id)), force=True)
    )


def get_document(
    *,
    corpus_id: str,
    document_id: str,
    client: genai.RetrieverServiceClient,
) -> Document | None:
    try:
        document = client.get_document(
            genai.GetDocumentRequest(
                name=str(EntityName(corpus_id=corpus_id, document_id=document_id))
            )
        )
        return Document.from_document(document)
    except Exception as e:
        if not isinstance(e, gapi_exception.NotFound):
            raise
        _logger.warning(f"Document {document_id} in corpus {corpus_id} not found: {e}")
        return None


def create_document(
    *,
    corpus_id: str,
    document_id: str | None = None,
    display_name: str | None = None,
    metadata: dict[str, Any] | None = None,
    client: genai.RetrieverServiceClient,
) -> Document:
    name: str | None
    if document_id is not None:
        name = str(EntityName(corpus_id=corpus_id, document_id=document_id))
    else:
        name = None

    new_display_name = display_name or f"Untitled {datetime.datetime.now()}"
    new_metadatas = _convert_to_metadata(metadata) if metadata else None

    new_document = client.create_document(
        genai.CreateDocumentRequest(
            parent=str(EntityName(corpus_id=corpus_id)),
            document=genai.Document(
                name=name, display_name=new_display_name, custom_metadata=new_metadatas
            ),
        )
    )

    return Document.from_document(new_document)


def delete_document(
    *,
    corpus_id: str,
    document_id: str,
    client: genai.RetrieverServiceClient,
) -> None:
    client.delete_document(
        genai.DeleteDocumentRequest(
            name=str(EntityName(corpus_id=corpus_id, document_id=document_id)),
            force=True,
        )
    )


def batch_create_chunk(
    *,
    corpus_id: str,
    document_id: str,
    texts: list[str],
    metadatas: list[dict[str, Any]] | None = None,
    client: genai.RetrieverServiceClient,
) -> list[genai.Chunk]:
    if metadatas is None:
        metadatas = [{} for _ in texts]
    if len(texts) != len(metadatas):
        msg = (
            f"metadatas's length {len(metadatas)} "
            f"and texts's length {len(texts)} are mismatched"
        )
        raise ValueError(msg)

    doc_name = str(EntityName(corpus_id=corpus_id, document_id=document_id))

    created_chunks: list[genai.Chunk] = []

    batch_request = genai.BatchCreateChunksRequest(
        parent=doc_name,
        requests=[],
    )
    for text, metadata in zip(texts, metadatas, strict=False):
        batch_request.requests.append(
            genai.CreateChunkRequest(
                parent=doc_name,
                chunk=genai.Chunk(
                    data=genai.ChunkData(string_value=text),
                    custom_metadata=_convert_to_metadata(metadata),
                ),
            )
        )

        if len(batch_request.requests) >= _MAX_REQUEST_PER_CHUNK:
            response = client.batch_create_chunks(batch_request)
            created_chunks.extend(list(response.chunks))
            # Prepare a new batch for next round.
            batch_request = genai.BatchCreateChunksRequest(
                parent=doc_name,
                requests=[],
            )

    # Process left over.
    if len(batch_request.requests) > 0:
        response = client.batch_create_chunks(batch_request)
        created_chunks.extend(list(response.chunks))

    return created_chunks


def delete_chunk(
    *,
    corpus_id: str,
    document_id: str,
    chunk_id: str,
    client: genai.RetrieverServiceClient,
) -> None:
    client.delete_chunk(
        genai.DeleteChunkRequest(
            name=str(
                EntityName(
                    corpus_id=corpus_id, document_id=document_id, chunk_id=chunk_id
                )
            )
        )
    )


def query_corpus(
    *,
    corpus_id: str,
    query: str,
    k: int = 4,
    filter: dict[str, Any] | None = None,
    client: genai.RetrieverServiceClient,
) -> list[genai.RelevantChunk]:
    response = client.query_corpus(
        genai.QueryCorpusRequest(
            name=str(EntityName(corpus_id=corpus_id)),
            query=query,
            metadata_filters=_convert_filter(filter),
            results_count=k,
        )
    )
    return list(response.relevant_chunks)


def query_document(
    *,
    corpus_id: str,
    document_id: str,
    query: str,
    k: int = 4,
    filter: dict[str, Any] | None = None,
    client: genai.RetrieverServiceClient,
) -> list[genai.RelevantChunk]:
    response = client.query_document(
        genai.QueryDocumentRequest(
            name=str(EntityName(corpus_id=corpus_id, document_id=document_id)),
            query=query,
            metadata_filters=_convert_filter(filter),
            results_count=k,
        )
    )
    return list(response.relevant_chunks)


def _convert_to_metadata(metadata: dict[str, Any]) -> list[genai.CustomMetadata]:
    cs: list[genai.CustomMetadata] = []
    for key, value in metadata.items():
        if isinstance(value, str):
            c = genai.CustomMetadata(key=key, string_value=value)
        elif isinstance(value, (float, int)):
            c = genai.CustomMetadata(key=key, numeric_value=value)
        else:
            msg = f"Metadata value {value} is not supported"
            raise ValueError(msg)

        cs.append(c)
    return cs


def _convert_filter(fs: dict[str, Any] | None) -> list[genai.MetadataFilter]:
    if fs is None:
        return []
    assert isinstance(fs, dict)

    filters: list[genai.MetadataFilter] = []
    for key, value in fs.items():
        if isinstance(value, str):
            condition = genai.Condition(
                operation=genai.Condition.Operator.EQUAL, string_value=value
            )
        elif isinstance(value, (float, int)):
            condition = genai.Condition(
                operation=genai.Condition.Operator.EQUAL, numeric_value=value
            )
        else:
            msg = f"Filter value {value} is not supported"
            raise ValueError(msg)

        filters.append(genai.MetadataFilter(key=key, conditions=[condition]))

    return filters
