from collections import UserList
from datetime import datetime
from functools import cached_property

from anthropic import Anthropic as AnthropicSDK

from kwark.ai_services import AIService
from wizlib.ui import Emphasis


DEFAULT_MODEL = 'claude-haiku-4-5-20251001'
DEFAULT_MAX_TOKENS = 4096
DEFAULT_TEMPERATURE = 0
SYSTEM_PROMPT_TEMPLATE = """You are Kwark, a lightweight AI assistant that
performs specific tasks. You use the {model} model. Now is {now}.\n"""


class AnthropicMessagesBlock(UserList):

    @staticmethod
    def _content(content):
        if isinstance(content, str):
            return [{"type": "text", "text": content}]
        else:
            return [dict(content)]

    def user_says(self, text):
        """Record what the user said to send to the model"""
        self.data.append({
            "role": "user",
            "content": text
        })

    def model_says(self, content):
        """Record what the model said for future loops"""
        self.data.append({
            "role": "assistant",
            "content": self._content(content)
        })

    def tool_says(self, id, result):
        self.data.append({
            "role": "user",
            "content": [{
                "type": "tool_result",
                "tool_use_id": id,
                "content": result
            }]
        })

    def need_user_input(self):
        return self.data[-1]['role'] == 'assistant'


class AnthropicAIService(AIService):

    service_type = 'anthropic'

    def __init__(self, api_key=None, model=None):
        if api_key:
            self.client = AnthropicSDK(api_key=api_key)
        else:
            self.client = AnthropicSDK()
        # self.available_models = self._fetch_available_models()
        self.model = DEFAULT_MODEL if model is None else model
        self.system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
            model=self.model_name, now=datetime.now())
        # Just trying to work with tools
        self.toolset = AnthropicToolset()

    @cached_property
    def available_models(self):
        """Fetch list of available AnthropicSDK models"""
        try:
            response = self.client.models.list()
            models = []
            for model in response:
                models.append({
                    'id': model.id,
                    'display_name': getattr(model, 'display_name', model.id),
                    'created_at': getattr(model, 'created_at', None),
                })
            return models
        except Exception as e:
            print(f"Error fetching models: {e}")
            return []

    @cached_property
    def model_name(self):
        return next((m['display_name'] for m in self.available_models
                     if m['id'] == self.model), None)

    def _api_arguments(self, messages: AnthropicMessagesBlock) -> dict:
        return {
            'model': self.model,
            'max_tokens': DEFAULT_MAX_TOKENS,
            'temperature': DEFAULT_TEMPERATURE,
            'system': self.system_prompt,
            'tools': self.toolset.definitions,
            'messages': messages
        }

    def query(self, text, model=None):
        """Send a single prompt and return the response"""
        messages = AnthropicMessagesBlock()
        messages.user_says(text)
        message = self.client.messages.create(**self._api_arguments(messages))
        return message.content[0].text

    def _stream(self, messages: AnthropicMessagesBlock, ui):
        """Streams text from the API to the ui, and returns the entire message object"""
        with self.client.messages.stream(**self._api_arguments(messages)) as stream:
            for text in stream.text_stream:
                ui.send(text, Emphasis.GENERAL, newline=False, wrap=80)
                # response += text
            message = stream.get_final_message()
        ui.send("\n")
        return message

    def chat(self, ui, initial_message=None):
        """Interactive chat with conversation history using UI"""
        messages = AnthropicMessagesBlock()
        if initial_message:
            messages.user_says(initial_message)
            response = self._stream(messages, ui)
            messages.model_says(response)
        while True:
            try:
                user_input = ui.get_text("\nYou: ").strip()
                if not user_input:
                    continue
                if user_input.lower() in ['quit', 'exit', 'bye']:
                    ui.send("Goodbye!", Emphasis.INFO)
                    break
                ui.send("\nKwark: ", Emphasis.GENERAL)
                messages.user_says(user_input)
                message = self._stream(messages, ui)
                while not messages.need_user_input():
                    for content in message.content:
                        if content.type == 'text':
                            messages.model_says(content.text)
                        elif content.type == 'tool_use':
                            if content.name == 'exit':  # A specialized tool
                                return
                            else:
                                messages.model_says(content)
                                id = content.id
                                ui.send(
                                    f'Using tool {content.name}...', emphasis=Emphasis.PRINCIPAL)
                                result = self.toolset.run_tool(
                                    content.name, content.input)
                                messages.tool_says(id, result)
                                message = self._stream(messages, ui)
            except KeyboardInterrupt:
                ui.send("Goodbye!", Emphasis.INFO)
                break
            except EOFError:
                ui.send("Goodbye!", Emphasis.INFO)
                break


class AnthropicToolset:
    """Local tools, on the road to MCP tools"""

    @property
    def definitions(self):
        return [{
            "name": "get_current_time",
            "description": "Returns the current time",
            "input_schema": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }]

    def run_tool(self, tool_name, arguments):
        """Execute local tools synchronously."""
        if tool_name == "get_current_time":
            return datetime.now().astimezone().isoformat()
        return None
