import asyncio

import tiktoken
from loguru import logger
from openai import AsyncClient

from app.config import JINA_API_KEY
from llm_generation.config import DEEPSEEK_MAX_TOKEN_LENGTH
from llm_generation.models.base import BaseModel


class Jina(BaseModel):
    def __init__(self, model_name: str = "jina-deepsearch-v1"):
        super().__init__(model_name)

    async def generate_response(
        self, user_prompt: str, conversation: list = None, **kwargs
    ) -> str:
        # Truncate the user prompt to MAX_TOKEN_LENGTH tokens, use tiktoken for deepseek temporarily
        tokenizer = tiktoken.encoding_for_model("gpt-4o")
        user_prompt_tokens = tokenizer.encode(user_prompt)
        if len(user_prompt_tokens) > DEEPSEEK_MAX_TOKEN_LENGTH:
            user_prompt = tokenizer.decode(
                user_prompt_tokens[:DEEPSEEK_MAX_TOKEN_LENGTH]
            )

        openai_client = AsyncClient(
            api_key=JINA_API_KEY, base_url="https://deepsearch.jina.ai/v1/"
        )
        conversation = conversation or []

        response = await openai_client.chat.completions.create(
            model=self.model_name,
            messages=conversation + [{"role": "user", "content": user_prompt}],
            **kwargs
        )
        # Streaming response
        if "stream" in kwargs:
            content = ""
            if self.streaming_callback is None:
                logger.warning(
                    "No streaming callback is set, skipping callback function"
                )

            async for chunk in response:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta

                # Call the streaming callback function
                if self.streaming_callback:
                    if asyncio.iscoroutinefunction(self.streaming_callback):
                        await self.streaming_callback(content, delta)
                    else:
                        self.streaming_callback(content, delta)

                # Append the content
                if delta.content:
                    content += delta.content
        else:
            content = response.choices[0].message.content
        return content


async def main():
    jina = Jina()
    jina.set_streaming_callback(print)
    conversation_history = [
        {
            "role": "system",
            "content": """You are an advanced AI research agent from Jina AI. You are specialized in multistep reasoning. Using your training data and prior lessons learned, answer the user question with absolute certainty.

Based on the current context, you must choose one of the following actions:
<actions>

<action-visit>
- Access and read full content from URLs
- Must check URLs mentioned in <question>

</action-visit>

<action-coding>
- This JavaScript-based solution helps you handle programming tasks like counting, filtering, transforming, sorting, regex extraction, and data processing.
- Simply describe your problem in the "codingIssue" field. Include actual values for small inputs or variable names for larger datasets.
- No code writing is required – senior engineers will handle the implementation.
</action-coding>

<action-search>
- Use web search to find relevant information
- Choose optimal search queries and language based on the expected answer format
- Focus on one specific aspect of the original question
- Suggest unique keywords and alternative search angles

</action-search>

<action-answer>
- For greetings, casual conversation, or general knowledge questions, answer directly without references.
- For all other questions, provide a verified answer with references. Each reference must include exactQuote and url.
- If uncertain, use <action-reflect>
</action-answer>

<action-reflect>
- Analyze through scenarios and systematic breakdowns
- Identify gaps and ask key clarifying questions that related to the original question and lead to the answer
</action-reflect>

</actions>
        """,
        },
    ]
    response = await jina.generate_response(
        conversation=conversation_history,
        user_prompt="what happened at crypto twitter today",
        stream=True,
    )
    print(response)


if __name__ == "__main__":
    asyncio.run(main())