from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
from google import genai
from google.genai import types
from pydantic import BaseModel

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from haema import EmbeddingClient, LLMClient, Memory

STORAGE_DIR = Path("./example_haema")
EMBEDDING_MODEL = "gemini-embedding-001"
LLM_MODEL = "gemini-2.5-flash"
OUTPUT_DIMENSIONALITY = 768

test = {
    "embed_query": 0,
    "embed_document": 0,
    "generate_structured": 0,
}


class GoogleGenAIEmbeddingClient(EmbeddingClient):
    def __init__(self, client: genai.Client, model: str) -> None:
        self.client = client
        self.model = model

    def embed_query(self, texts: list[str], output_dimensionality: int) -> np.ndarray:
        test["embed_query"] += 1
        return self._embed(
            texts=texts,
            output_dimensionality=output_dimensionality,
            task_type="RETRIEVAL_QUERY",
        )

    def embed_document(self, texts: list[str], output_dimensionality: int) -> np.ndarray:
        test["embed_document"] += 1
        return self._embed(
            texts=texts,
            output_dimensionality=output_dimensionality,
            task_type="RETRIEVAL_DOCUMENT",
        )

    def _embed(
        self,
        texts: list[str],
        output_dimensionality: int,
        task_type: str,
    ) -> np.ndarray:
        response = self.client.models.embed_content(
            model=self.model,
            contents=texts,
            config=types.EmbedContentConfig(
                output_dimensionality=output_dimensionality,
                task_type=task_type,
            ),
        )

        vectors = [item.values for item in (response.embeddings or [])]
        return np.asarray(vectors, dtype=np.float32)


class GoogleGenAILLMClient(LLMClient):
    def __init__(self, client: genai.Client, model: str) -> None:
        self.client = client
        self.model = model

    def generate_structured(
        self,
        system_prompt: str,
        user_prompt: str,
        response_model: type[BaseModel],
    ) -> dict[str, Any]:
        test["generate_structured"] += 1
        response = self.client.models.generate_content(
            model=self.model,
            contents=user_prompt,
            config=types.GenerateContentConfig(
                system_instruction=system_prompt,
                response_mime_type="application/json",
                response_schema=response_model.model_json_schema(),
            ),
        )

        if isinstance(response.parsed, dict):
            return response.parsed
        if isinstance(response.text, str) and response.text.strip():
            return json.loads(response.text)
        raise ValueError("LLM response did not contain parseable JSON")


def main() -> None:
    api_key = os.getenv("GOOGLE_API_KEY", "").strip()
    if not api_key:
        raise SystemExit("Missing GOOGLE_API_KEY environment variable.")

    client = genai.Client(api_key=api_key)
    embedding_client = GoogleGenAIEmbeddingClient(client=client, model=EMBEDDING_MODEL)
    llm_client = GoogleGenAILLMClient(client=client, model=LLM_MODEL)

    memory = Memory(
        path=STORAGE_DIR,
        output_dimensionality=OUTPUT_DIMENSIONALITY,
        embedding_client=embedding_client,
        llm_client=llm_client,
    )

    text: str = (
        "User prefers concise, actionable answers."
        "\n"
        "User works on HAEMA memory framework using ChromaDB."
        "\n"
        "User asked to test google-genai with Gemini flash and embedding models."
        "\n"
        "User loves a cat."
    )

    memory.add(text)

    print("=== CORE ===")
    print(memory.get_core())
    print("=== LATEST(1,3) ===")
    print(memory.get_latest(begin=1, count=3))
    print("=== SEARCH ===")
    print(memory.search("What model does the user want?", n=3))
    print("=== TEST COUNTS ===")
    print(test)


if __name__ == "__main__":
    import time
    start = time.time()
    main()
    print(f"Execution time: {time.time() - start:.2f} seconds")
