import uuid
import time
from datetime import datetime
from typing import Any, Optional, Iterator
from google import genai
from .storage import LocalStorage, IStorage
from .remote_storage import RemoteStorage
from .types import TrackedRequest, TrackerOptions


class Client:
    def __init__(
        self,
        api_key: Optional[str] = None,
        vertexai: bool = False,
        project: Optional[str] = None,
        location: Optional[str] = None,
        storage: str = "local",
        local_storage: Optional[dict] = None,
        remote_storage: Optional[dict] = None,
        **kwargs
    ):
        storage_mode = storage
        local_storage_options = local_storage
        remote_storage_options = remote_storage

        client_kwargs = {}
        if api_key:
            client_kwargs["api_key"] = api_key
        if vertexai:
            client_kwargs["vertexai"] = True
            if project:
                client_kwargs["project"] = project
            if location:
                client_kwargs["location"] = location

        self._client = genai.Client(**client_kwargs)

        if storage_mode == "remote":
            if not remote_storage_options:
                raise ValueError("remote_storage options required when storage mode is 'remote'")
            self._storage: IStorage = RemoteStorage(remote_storage_options)
        else:
            self._storage: IStorage = LocalStorage(local_storage_options)
        self._original_generate_content = self._client.models.generate_content
        self._original_generate_content_stream = self._client.models.generate_content_stream

        self._client.models.generate_content = self._wrapped_generate_content
        self._client.models.generate_content_stream = self._wrapped_generate_content_stream

    def _wrapped_generate_content(self, **params: Any) -> Any:
        request_id = str(uuid.uuid4())
        timestamp = datetime.now()
        start_time = time.perf_counter()

        metadata = params.pop("metadata", None)
        model = params.get("model", "unknown")
        contents = params.get("contents")

        config = params.get("config")
        tools = config.tools if config and hasattr(config, "tools") else params.get("tools")
        search_grounding = bool(
            tools and hasattr(tools, "__iter__") and
            any(hasattr(tool, "google_search") for tool in tools)
        )

        try:
            response = self._original_generate_content(**params)
            response_time = (time.perf_counter() - start_time) * 1000

            received_search_grounded_response = bool(
                hasattr(response, "candidates") and
                response.candidates and
                hasattr(response.candidates[0], "grounding_metadata") and
                response.candidates[0].grounding_metadata and
                (hasattr(response.candidates[0].grounding_metadata, "search_entry_point") or
                 hasattr(response.candidates[0].grounding_metadata, "grounding_chunks"))
            )

            usage_metadata = response.usage_metadata if hasattr(response, "usage_metadata") else None

            response_dict = self._serialize_response(response)

            tracked_request: TrackedRequest = {
                "id": request_id,
                "timestamp": timestamp,
                "response_time": response_time,
                "model": model,
                "prompt": contents,
                "response": response_dict,
                "usage": {
                    "promptTokens": usage_metadata.prompt_token_count if usage_metadata else None,
                    "completionTokens": usage_metadata.candidates_token_count if usage_metadata else None,
                    "totalTokens": usage_metadata.total_token_count if usage_metadata else None,
                },
                "search_grounding": search_grounding,
                "received_search_grounded_response": received_search_grounded_response,
                "metadata": metadata,
            }

            try:
                self._storage.add(tracked_request)
            except Exception as storage_error:
                print(f"Failed to save tracked request: {storage_error}")

            return response

        except Exception as error:
            tracked_request: TrackedRequest = {
                "id": request_id,
                "timestamp": timestamp,
                "model": model,
                "prompt": contents,
                "response": None,
                "error": str(error),
                "search_grounding": search_grounding,
                "received_search_grounded_response": False,
                "metadata": metadata,
            }

            try:
                self._storage.add(tracked_request)
            except Exception as storage_error:
                print(f"Failed to save tracked request with error: {storage_error}")

            raise

    def _serialize_response(self, response: Any) -> dict:
        result = {}

        if hasattr(response, "sdk_http_response"):
            sdk_resp = response.sdk_http_response
            result["sdkHttpResponse"] = {
                "headers": dict(sdk_resp.headers) if hasattr(sdk_resp, "headers") else {}
            }

        if hasattr(response, "candidates"):
            result["candidates"] = []
            for candidate in response.candidates:
                cand_dict = {}
                if hasattr(candidate, "content"):
                    content = candidate.content
                    cand_dict["content"] = {
                        "parts": [{"text": part.text} for part in content.parts if hasattr(part, "text")],
                        "role": content.role if hasattr(content, "role") else None
                    }
                if hasattr(candidate, "finish_reason"):
                    cand_dict["finishReason"] = str(candidate.finish_reason).replace("FinishReason.", "")
                if hasattr(candidate, "grounding_metadata"):
                    cand_dict["groundingMetadata"] = {}
                if hasattr(candidate, "index"):
                    cand_dict["index"] = candidate.index
                result["candidates"].append(cand_dict)

        if hasattr(response, "model_version"):
            result["modelVersion"] = response.model_version
        if hasattr(response, "response_id"):
            result["responseId"] = response.response_id

        if hasattr(response, "usage_metadata"):
            um = response.usage_metadata
            result["usageMetadata"] = {
                "promptTokenCount": getattr(um, "prompt_token_count", None),
                "candidatesTokenCount": getattr(um, "candidates_token_count", None),
                "totalTokenCount": getattr(um, "total_token_count", None),
            }
            if hasattr(um, "prompt_tokens_details"):
                result["usageMetadata"]["promptTokensDetails"] = [
                    {"modality": str(d.modality).replace("MediaModality.", ""), "tokenCount": d.token_count}
                    for d in um.prompt_tokens_details
                ]
            if hasattr(um, "thoughts_token_count"):
                result["usageMetadata"]["thoughtsTokenCount"] = um.thoughts_token_count

        return result

    def _wrapped_generate_content_stream(self, **params: Any) -> Iterator[Any]:
        request_id = str(uuid.uuid4())
        timestamp = datetime.now()
        start_time = time.perf_counter()

        metadata = params.pop("metadata", None)
        model = params.get("model", "unknown")
        contents = params.get("contents")

        config = params.get("config")
        tools = config.tools if config and hasattr(config, "tools") else params.get("tools")
        search_grounding = bool(
            tools and hasattr(tools, "__iter__") and
            any(hasattr(tool, "google_search") for tool in tools)
        )

        try:
            stream = self._original_generate_content_stream(**params)
            chunks = []
            full_response = None

            for chunk in stream:
                chunks.append(chunk)
                yield chunk

            if chunks:
                full_response = chunks[-1]

            response_time = (time.perf_counter() - start_time) * 1000

            received_search_grounded_response = bool(
                full_response and
                hasattr(full_response, "candidates") and
                full_response.candidates and
                hasattr(full_response.candidates[0], "grounding_metadata") and
                full_response.candidates[0].grounding_metadata and
                (hasattr(full_response.candidates[0].grounding_metadata, "search_entry_point") or
                 hasattr(full_response.candidates[0].grounding_metadata, "grounding_chunks"))
            )

            usage_metadata = full_response.usage_metadata if full_response and hasattr(full_response, "usage_metadata") else None

            response_dict = self._serialize_response(full_response) if full_response else None

            tracked_request: TrackedRequest = {
                "id": request_id,
                "timestamp": timestamp,
                "response_time": response_time,
                "model": model,
                "prompt": contents,
                "response": response_dict,
                "usage": {
                    "promptTokens": usage_metadata.prompt_token_count if usage_metadata else None,
                    "completionTokens": usage_metadata.candidates_token_count if usage_metadata else None,
                    "totalTokens": usage_metadata.total_token_count if usage_metadata else None,
                },
                "search_grounding": search_grounding,
                "received_search_grounded_response": received_search_grounded_response,
                "metadata": metadata,
            }

            try:
                self._storage.add(tracked_request)
            except Exception as storage_error:
                print(f"Failed to save streamed tracked request: {storage_error}")

        except Exception as error:
            tracked_request: TrackedRequest = {
                "id": request_id,
                "timestamp": timestamp,
                "model": model,
                "prompt": contents,
                "response": None,
                "error": str(error),
                "search_grounding": search_grounding,
                "received_search_grounded_response": False,
                "metadata": metadata,
            }

            try:
                self._storage.add(tracked_request)
            except Exception as storage_error:
                print(f"Failed to save streamed tracked request with error: {storage_error}")

            raise

    @property
    def models(self):
        return self._client.models

    @property
    def chats(self):
        return self._client.chats

    @property
    def files(self):
        return self._client.files

    @property
    def caches(self):
        return self._client.caches

    def get_tracked_requests(self) -> list[TrackedRequest]:
        return self._storage.get_all()

    def load_history(self) -> None:
        if isinstance(self._storage, LocalStorage):
            self._storage.load()

    def save_history(self) -> None:
        if isinstance(self._storage, LocalStorage):
            self._storage.save()

    def clear_history(self) -> None:
        self._storage.clear()
