import json
from typing import AsyncIterator

from anthropic import AsyncAnthropic, AsyncAnthropicVertex
from loguru import logger

from .model_abstract import (
    ContentType,
    FunctionCall,
    GenerateRequest,
    GenerateResponse,
    LLMModelAbstract,
    Message,
    ModelCapability,
    StreamChunk,
    Tool,
    ToolCall,
)


class ClaudeModel(LLMModelAbstract):
    """Anthropic Claude model implementation."""

    name = "claude"

    def __init__(
        self,
        model_name: str,
        api_key: str,
        base_url: str | None = None,
    ):
        """
        Initialize Claude model via Anthropic API.

        Args:
            model_name: Model identifier (e.g., "claude-3-5-sonnet-20241022", "claude-3-opus-20240229")
            api_key: Anthropic API key
            base_url: Optional custom base URL
        """
        self._model_name = model_name
        self.client = AsyncAnthropic(
            api_key=api_key,
            base_url=base_url,
        )
        self._capabilities = self._determine_capabilities()

    def _determine_capabilities(self) -> ModelCapability:
        """Determine capabilities based on model name."""
        caps = (
            ModelCapability.TEXT_GENERATION
            | ModelCapability.STREAMING
            | ModelCapability.TOOL_CALLING
            | ModelCapability.STRUCTURED_OUTPUT
            | ModelCapability.VISION
            | ModelCapability.MULTIMODAL_INPUT
        )
        return caps

    @property
    def model_name(self) -> str:
        return self._model_name

    @model_name.setter
    def model_name(self, value: str) -> None:
        self._model_name = value
        self._capabilities = self._determine_capabilities()

    @property
    def capabilities(self) -> ModelCapability:
        return self._capabilities

    async def list_chat_models(self) -> list[str]:
        """List available Claude models from Anthropic API."""
        try:
            response = await self.client.models.list()
            # Extract model IDs from the response
            models = [model.id for model in response.data]
            return sorted(models)
        except Exception as e:
            logger.warning("Failed to list Anthropic models: {}", e)
            return []

    def _convert_message(self, msg: Message) -> dict:
        """Convert internal Message to Claude format."""
        # Tool result messages: Claude expects role=user with tool_result content blocks
        if msg.role == "tool":
            return {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": msg.tool_call_id,
                        "content": msg.content
                        if isinstance(msg.content, str)
                        else json.dumps(msg.content),
                    }
                ],
            }

        # Assistant messages with tool_calls: Claude expects tool_use content blocks
        if msg.role == "assistant" and msg.tool_calls:
            content_parts = []
            if msg.content:
                text = msg.content if isinstance(msg.content, str) else ""
                if text:
                    content_parts.append({"type": "text", "text": text})
            for tc in msg.tool_calls:
                content_parts.append(
                    {
                        "type": "tool_use",
                        "id": tc.id,
                        "name": tc.function.name,
                        "input": json.loads(tc.function.arguments)
                        if isinstance(tc.function.arguments, str)
                        else tc.function.arguments,
                    }
                )
            return {"role": "assistant", "content": content_parts}

        result = {"role": msg.role}

        # Claude uses "user" and "assistant" roles, no "system" in messages array
        if msg.role == "system":
            result["role"] = "user"

        # Handle content
        if isinstance(msg.content, str):
            result["content"] = msg.content
        elif msg.content is not None:
            # Multimodal content
            content_parts = []
            for part in msg.content:
                if part.content_type == ContentType.TEXT:
                    content_parts.append({"type": "text", "text": part.content})
                elif part.content_type == ContentType.IMAGE_URL:
                    content_parts.append(
                        {
                            "type": "image",
                            "source": {
                                "type": "url",
                                "url": part.content,
                            },
                        }
                    )
                elif part.content_type == ContentType.IMAGE_BASE64:
                    content_parts.append(
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": part.mime_type or "image/jpeg",
                                "data": part.content,
                            },
                        }
                    )
            result["content"] = content_parts
        else:
            result["content"] = ""

        return result

    STRUCTURED_OUTPUT_TOOL_NAME = "structured_output"

    def _convert_tools(self, tools: list[Tool]) -> list[dict]:
        """Convert internal Tool definitions to Claude format."""
        return [
            {
                "name": tool.function.name,
                "description": tool.function.description,
                "input_schema": tool.function.parameters,
            }
            for tool in tools
        ]

    def _apply_structured_output(self, kwargs: dict, response_format: dict) -> bool:
        """Apply response_format as a forced tool use for structured output.

        Returns True if structured output was applied.
        """
        schema = response_format
        if "schema" in response_format:
            schema = response_format["schema"]

        tool = {
            "name": self.STRUCTURED_OUTPUT_TOOL_NAME,
            "description": "Return the structured output in the requested format.",
            "input_schema": schema,
        }
        tools = kwargs.get("tools", [])
        tools.append(tool)
        kwargs["tools"] = tools
        kwargs["tool_choice"] = {
            "type": "tool",
            "name": self.STRUCTURED_OUTPUT_TOOL_NAME,
        }
        return True

    def _extract_structured_response(self, response) -> GenerateResponse:
        """Extract structured output from a tool_use response."""
        for block in response.content:
            if (
                block.type == "tool_use"
                and block.name == self.STRUCTURED_OUTPUT_TOOL_NAME
            ):
                return GenerateResponse(
                    content=json.dumps(block.input),
                    tool_calls=None,
                    finish_reason=response.stop_reason,
                    usage={
                        "prompt_tokens": response.usage.input_tokens,
                        "completion_tokens": response.usage.output_tokens,
                        "total_tokens": response.usage.input_tokens
                        + response.usage.output_tokens,
                    },
                )
        # Fallback: no structured_output tool found
        return self._extract_standard_response(response)

    def _extract_standard_response(self, response) -> GenerateResponse:
        """Extract standard response (text + tool calls)."""
        content = None
        text_blocks = [block.text for block in response.content if block.type == "text"]
        if text_blocks:
            content = "".join(text_blocks)

        tool_calls = None
        tool_use_blocks = [
            block for block in response.content if block.type == "tool_use"
        ]
        if tool_use_blocks:
            tool_calls = [
                ToolCall(
                    id=block.id,
                    type="function",
                    function=FunctionCall(
                        name=block.name,
                        arguments=json.dumps(block.input),
                    ),
                )
                for block in tool_use_blocks
            ]

        return GenerateResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason=response.stop_reason,
            usage={
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens
                + response.usage.output_tokens,
            },
        )

    async def generate(self, request: GenerateRequest) -> GenerateResponse:
        """Generate a response using Claude API."""
        await self.validate_request(request)

        # Extract system message
        system_message = None
        messages = []
        for msg in request.messages:
            if msg.role == "system":
                system_message = msg.content if isinstance(msg.content, str) else ""
            else:
                messages.append(self._convert_message(msg))

        kwargs = {
            "model": self._model_name,
            "messages": messages,
            "max_tokens": request.max_tokens or 4096,  # Claude requires max_tokens
        }

        if system_message:
            kwargs["system"] = system_message
        if request.temperature is not None:
            kwargs["temperature"] = request.temperature
        if request.top_p is not None:
            kwargs["top_p"] = request.top_p
        if request.stop:
            kwargs["stop_sequences"] = request.stop
        if request.tools:
            kwargs["tools"] = self._convert_tools(request.tools)

        use_structured = False
        if request.response_format:
            use_structured = self._apply_structured_output(
                kwargs, request.response_format
            )

        response = await self.client.messages.create(**kwargs)

        if use_structured:
            return self._extract_structured_response(response)
        return self._extract_standard_response(response)

    async def generate_stream(
        self, request: GenerateRequest
    ) -> AsyncIterator[StreamChunk]:
        """Generate a streaming response using Claude API."""
        await self.validate_request(request)

        # Extract system message
        system_message = None
        messages = []
        for msg in request.messages:
            if msg.role == "system":
                system_message = msg.content if isinstance(msg.content, str) else ""
            else:
                messages.append(self._convert_message(msg))

        kwargs = {
            "model": self._model_name,
            "messages": messages,
            "max_tokens": request.max_tokens or 4096,
        }

        if system_message:
            kwargs["system"] = system_message
        if request.temperature is not None:
            kwargs["temperature"] = request.temperature
        if request.top_p is not None:
            kwargs["top_p"] = request.top_p
        if request.stop:
            kwargs["stop_sequences"] = request.stop
        if request.tools:
            kwargs["tools"] = self._convert_tools(request.tools)

        async with self.client.messages.stream(**kwargs) as stream:
            async for event in stream:
                content = None
                tool_calls = None
                finish_reason = None

                if event.type == "content_block_delta":
                    if hasattr(event.delta, "text"):
                        content = event.delta.text

                elif event.type == "content_block_stop":
                    if (
                        hasattr(event, "content_block")
                        and event.content_block.type == "tool_use"
                    ):
                        tool_calls = [
                            ToolCall(
                                id=event.content_block.id,
                                type="function",
                                function=FunctionCall(
                                    name=event.content_block.name,
                                    arguments=json.dumps(event.content_block.input),
                                ),
                            )
                        ]

                elif event.type == "message_stop":
                    finish_reason = "stop"

                if content or tool_calls or finish_reason:
                    yield StreamChunk(
                        content=content,
                        tool_calls=tool_calls,
                        finish_reason=finish_reason,
                    )


class ClaudeVertexModel(LLMModelAbstract):
    """Anthropic Claude model via Vertex AI."""

    def __init__(
        self,
        model_name: str,
        project_id: str,
        location: str = "us-east5",
        credentials: dict | None = None,
    ):
        """
        Initialize Claude model via Vertex AI.

        Args:
            model_name: Model identifier (e.g., "claude-3-5-sonnet-v2@20241022")
            project_id: GCP project ID
            location: GCP location (us-east5 for Claude)
            credentials: Optional service account credentials dict (JSON key contents)
        """
        self._model_name = model_name
        self._project_id = project_id
        self._location = location

        kwargs = {
            "project_id": project_id,
            "region": location,
        }

        if credentials:
            from google.oauth2 import service_account

            creds = service_account.Credentials.from_service_account_info(
                credentials,
                scopes=["https://www.googleapis.com/auth/cloud-platform"],
            )
            kwargs["credentials"] = creds

        self.client = AsyncAnthropicVertex(**kwargs)
        self._capabilities = self._determine_capabilities()

    def _determine_capabilities(self) -> ModelCapability:
        """Determine capabilities based on model name."""
        caps = (
            ModelCapability.TEXT_GENERATION
            | ModelCapability.STREAMING
            | ModelCapability.TOOL_CALLING
            | ModelCapability.STRUCTURED_OUTPUT
            | ModelCapability.VISION
            | ModelCapability.MULTIMODAL_INPUT
        )
        return caps

    @property
    def model_name(self) -> str:
        return self._model_name

    @model_name.setter
    def model_name(self, value: str) -> None:
        self._model_name = value
        self._capabilities = self._determine_capabilities()

    @property
    def capabilities(self) -> ModelCapability:
        return self._capabilities

    async def list_chat_models(self) -> list[str]:
        """List available Claude models via Vertex AI."""
        # Vertex AI doesn't provide a list models API for Claude
        # Return commonly available Claude models on Vertex AI
        from loguru import logger

        logger.warning(
            "Vertex AI does not support listing Claude models via API. "
            "Returning known available models."
        )
        # Common Claude models available on Vertex AI
        return [
            "claude-3-5-sonnet-v2@20241022",
            "claude-3-5-sonnet@20240620",
            "claude-3-opus@20240229",
            "claude-3-sonnet@20240229",
            "claude-3-haiku@20240307",
        ]

    def _convert_message(self, msg: Message) -> dict:
        """Convert internal Message to Claude format."""
        # Tool result messages: Claude expects role=user with tool_result content blocks
        if msg.role == "tool":
            return {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": msg.tool_call_id,
                        "content": msg.content
                        if isinstance(msg.content, str)
                        else json.dumps(msg.content),
                    }
                ],
            }

        # Assistant messages with tool_calls: Claude expects tool_use content blocks
        if msg.role == "assistant" and msg.tool_calls:
            content_parts = []
            if msg.content:
                text = msg.content if isinstance(msg.content, str) else ""
                if text:
                    content_parts.append({"type": "text", "text": text})
            for tc in msg.tool_calls:
                content_parts.append(
                    {
                        "type": "tool_use",
                        "id": tc.id,
                        "name": tc.function.name,
                        "input": json.loads(tc.function.arguments)
                        if isinstance(tc.function.arguments, str)
                        else tc.function.arguments,
                    }
                )
            return {"role": "assistant", "content": content_parts}

        result = {"role": msg.role}

        if msg.role == "system":
            result["role"] = "user"

        # Handle content
        if isinstance(msg.content, str):
            result["content"] = msg.content
        elif msg.content is not None:
            # Multimodal content
            content_parts = []
            for part in msg.content:
                if part.content_type == ContentType.TEXT:
                    content_parts.append({"type": "text", "text": part.content})
                elif part.content_type == ContentType.IMAGE_URL:
                    content_parts.append(
                        {
                            "type": "image",
                            "source": {
                                "type": "url",
                                "url": part.content,
                            },
                        }
                    )
                elif part.content_type == ContentType.IMAGE_BASE64:
                    content_parts.append(
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": part.mime_type or "image/jpeg",
                                "data": part.content,
                            },
                        }
                    )
            result["content"] = content_parts
        else:
            result["content"] = ""

        return result

    STRUCTURED_OUTPUT_TOOL_NAME = "structured_output"

    def _convert_tools(self, tools: list[Tool]) -> list[dict]:
        """Convert internal Tool definitions to Claude format."""
        return [
            {
                "name": tool.function.name,
                "description": tool.function.description,
                "input_schema": tool.function.parameters,
            }
            for tool in tools
        ]

    def _apply_structured_output(self, kwargs: dict, response_format: dict) -> bool:
        """Apply response_format as a forced tool use for structured output."""
        schema = response_format
        if "schema" in response_format:
            schema = response_format["schema"]

        tool = {
            "name": self.STRUCTURED_OUTPUT_TOOL_NAME,
            "description": "Return the structured output in the requested format.",
            "input_schema": schema,
        }
        tools = kwargs.get("tools", [])
        tools.append(tool)
        kwargs["tools"] = tools
        kwargs["tool_choice"] = {
            "type": "tool",
            "name": self.STRUCTURED_OUTPUT_TOOL_NAME,
        }
        return True

    def _extract_structured_response(self, response) -> GenerateResponse:
        """Extract structured output from a tool_use response."""
        for block in response.content:
            if (
                block.type == "tool_use"
                and block.name == self.STRUCTURED_OUTPUT_TOOL_NAME
            ):
                return GenerateResponse(
                    content=json.dumps(block.input),
                    tool_calls=None,
                    finish_reason=response.stop_reason,
                    usage={
                        "prompt_tokens": response.usage.input_tokens,
                        "completion_tokens": response.usage.output_tokens,
                        "total_tokens": response.usage.input_tokens
                        + response.usage.output_tokens,
                    },
                )
        return self._extract_standard_response(response)

    def _extract_standard_response(self, response) -> GenerateResponse:
        """Extract standard response (text + tool calls)."""
        content = None
        text_blocks = [block.text for block in response.content if block.type == "text"]
        if text_blocks:
            content = "".join(text_blocks)

        tool_calls = None
        tool_use_blocks = [
            block for block in response.content if block.type == "tool_use"
        ]
        if tool_use_blocks:
            tool_calls = [
                ToolCall(
                    id=block.id,
                    type="function",
                    function=FunctionCall(
                        name=block.name,
                        arguments=json.dumps(block.input),
                    ),
                )
                for block in tool_use_blocks
            ]

        return GenerateResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason=response.stop_reason,
            usage={
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens
                + response.usage.output_tokens,
            },
        )

    async def generate(self, request: GenerateRequest) -> GenerateResponse:
        """Generate a response using Claude via Vertex AI."""
        await self.validate_request(request)

        # Extract system message
        system_message = None
        messages = []
        for msg in request.messages:
            if msg.role == "system":
                system_message = msg.content if isinstance(msg.content, str) else ""
            else:
                messages.append(self._convert_message(msg))

        kwargs = {
            "model": self._model_name,
            "messages": messages,
            "max_tokens": request.max_tokens or 4096,
        }

        if system_message:
            kwargs["system"] = system_message
        if request.temperature is not None:
            kwargs["temperature"] = request.temperature
        if request.top_p is not None:
            kwargs["top_p"] = request.top_p
        if request.stop:
            kwargs["stop_sequences"] = request.stop
        if request.tools:
            kwargs["tools"] = self._convert_tools(request.tools)

        use_structured = False
        if request.response_format:
            use_structured = self._apply_structured_output(
                kwargs, request.response_format
            )

        response = await self.client.messages.create(**kwargs)

        if use_structured:
            return self._extract_structured_response(response)
        return self._extract_standard_response(response)

    async def generate_stream(
        self, request: GenerateRequest
    ) -> AsyncIterator[StreamChunk]:
        """Generate a streaming response using Claude via Vertex AI."""
        await self.validate_request(request)

        # Extract system message
        system_message = None
        messages = []
        for msg in request.messages:
            if msg.role == "system":
                system_message = msg.content if isinstance(msg.content, str) else ""
            else:
                messages.append(self._convert_message(msg))

        kwargs = {
            "model": self._model_name,
            "messages": messages,
            "max_tokens": request.max_tokens or 4096,
        }

        if system_message:
            kwargs["system"] = system_message
        if request.temperature is not None:
            kwargs["temperature"] = request.temperature
        if request.top_p is not None:
            kwargs["top_p"] = request.top_p
        if request.stop:
            kwargs["stop_sequences"] = request.stop
        if request.tools:
            kwargs["tools"] = self._convert_tools(request.tools)

        async with self.client.messages.stream(**kwargs) as stream:
            async for event in stream:
                content = None
                tool_calls = None
                finish_reason = None

                if event.type == "content_block_delta":
                    if hasattr(event.delta, "text"):
                        content = event.delta.text

                elif event.type == "content_block_stop":
                    if (
                        hasattr(event, "content_block")
                        and event.content_block.type == "tool_use"
                    ):
                        tool_calls = [
                            ToolCall(
                                id=event.content_block.id,
                                type="function",
                                function=FunctionCall(
                                    name=event.content_block.name,
                                    arguments=json.dumps(event.content_block.input),
                                ),
                            )
                        ]

                elif event.type == "message_stop":
                    finish_reason = "stop"

                if content or tool_calls or finish_reason:
                    yield StreamChunk(
                        content=content,
                        tool_calls=tool_calls,
                        finish_reason=finish_reason,
                    )
