import base64
import json
import logging
import os
import time
from dataclasses import dataclass
from functools import cached_property
from typing import Optional

from openai import OpenAI, AsyncOpenAI

from .converse import (Converse, ConverseAgent, StructuredConverse, ConverseResponse, ConverseOutput,
                       Message, MessageContent, ToolUse, TokenUsage, ConverseMetrics,
                       ReasoningContent, ReasoningText, AdditionalModelRequestFields, ThinkingConfig)

logger = logging.getLogger(__name__)

STOP_REASON_MAP = {'stop': 'end_turn', 'length': 'max_tokens', 'tool_calls': 'tool_use'}


class _MantleTransport:
    api_key: Optional[str] = None

    @property
    def _mantle_base_url(self):
        if endpoint := os.environ.get('MANTLE_ENDPOINT'):
            return endpoint
        region = self.region_name or self.session.region_name
        return f'https://bedrock-mantle.{region}.api.aws/v1'

    def _get_client(self, openai_class):
        if self.api_key:
            return openai_class(api_key=self.api_key, base_url=self._mantle_base_url)
        if api_key := os.environ.get('MANTLE_API_KEY'):
            return openai_class(api_key=api_key, base_url=self._mantle_base_url)
        return openai_class(base_url=self._mantle_base_url)

    @cached_property
    def openai_client(self) -> OpenAI:
        return self._get_client(OpenAI)

    @cached_property
    def async_openai_client(self) -> AsyncOpenAI:
        return self._get_client(AsyncOpenAI)

    def _build_tool_params(self, params):
        if not self.tool_config:
            return
        tools = [{'type': 'function', 'function': {
            'name': t.tool_spec.name, 'description': t.tool_spec.description,
            'parameters': t.tool_spec.input_schema.get('json', {})
        }} for t in self.tool_config.tools if t.tool_spec]
        if tools:
            params['tools'] = tools
        if self.tool_config.tool_choice:
            tc = self.tool_config.tool_choice
            if tc.tool:
                params['tool_choice'] = {'type': 'function', 'function': {'name': tc.tool.name}}
            elif tc.any:
                params['tool_choice'] = 'required'
            elif tc.auto:
                params['tool_choice'] = 'auto'

    def _build_inference_params(self, params):
        if not self.inference_config:
            return
        ic = self.inference_config
        if ic.max_tokens is not None: params['max_tokens'] = ic.max_tokens
        if ic.temperature is not None: params['temperature'] = ic.temperature
        if ic.top_p is not None: params['top_p'] = ic.top_p
        if ic.stop_sequences: params['stop'] = ic.stop_sequences

    def _build_thinking_params(self, params):
        if not (self.additional_model_request_fields and self.additional_model_request_fields.thinking
                and self.additional_model_request_fields.thinking.type == 'enabled'):
            return
        budget = self.additional_model_request_fields.thinking.budget_tokens
        if isinstance(budget, str):
            params['reasoning_effort'] = budget
        else:
            if budget <= 2048: params['reasoning_effort'] = 'low'
            elif budget <= 8192: params['reasoning_effort'] = 'medium'
            else: params['reasoning_effort'] = 'high'

    def _build_params(self, messages=None) -> dict:
        self.remove_invalid_caching(messages)
        msgs = []
        if self.system:
            system_text = '\n'.join(s.text for s in self.system if s.text)
            if system_text:
                msgs.append({'role': 'system', 'content': system_text})
        for msg in (messages or self.messages):
            msgs.extend(self._convert_message(msg))
        params = {'model': self.model_id, 'messages': msgs}
        self._build_tool_params(params)
        self._build_inference_params(params)
        self._build_thinking_params(params)
        return params

    def _convert_tool_results(self, content_list):
        results = []
        for c in content_list:
            if not c.tool_result:
                continue
            tr = c.tool_result
            parts = []
            for trc in tr.content:
                if trc.text: parts.append(trc.text)
                elif trc.json is not None: parts.append(json.dumps(trc.json))
            results.append({'role': 'tool', 'tool_call_id': tr.tool_use_id, 'content': '\n'.join(parts)})
        return results

    def _convert_assistant(self, content_list):
        openai_msg = {'role': 'assistant'}
        texts, tool_calls = [], []
        for c in content_list:
            if c.text: texts.append(c.text)
            elif c.tool_use:
                tool_calls.append({'id': c.tool_use.tool_use_id, 'type': 'function',
                    'function': {'name': c.tool_use.name, 'arguments': json.dumps(c.tool_use.input) if isinstance(c.tool_use.input, dict) else str(c.tool_use.input)}})
        if texts: openai_msg['content'] = '\n'.join(texts)
        if tool_calls: openai_msg['tool_calls'] = tool_calls
        return [openai_msg]

    def _convert_user(self, content_list):
        parts, has_multimodal = [], False
        for c in content_list:
            if c.text: parts.append({'type': 'text', 'text': c.text})
            elif c.image:
                has_multimodal = True
                b64 = base64.b64encode(c.image.source.bytes).decode()
                parts.append({'type': 'image_url', 'image_url': {'url': f'data:image/{c.image.format};base64,{b64}'}})
            elif c.document:
                b64 = base64.b64encode(c.document.source.bytes).decode()
                parts.append({'type': 'text', 'text': f'[Document: {c.document.name}.{c.document.format}]\n{b64}'})
        if not parts:
            return []
        if has_multimodal or len(parts) > 1:
            return [{'role': 'user', 'content': parts}]
        return [{'role': 'user', 'content': parts[0].get('text', '')}]

    def _convert_message(self, msg):
        tool_results = self._convert_tool_results(msg.content)
        other = [c for c in msg.content if not c.tool_result and not c.cache_point]
        if not other:
            return tool_results
        if msg.role == 'assistant':
            return tool_results + self._convert_assistant(other)
        return tool_results + self._convert_user(other)

    def _parse_completion(self, completion, latency_ms) -> ConverseResponse:
        if not completion.choices:
            return ConverseResponse(
                output=ConverseOutput(message=Message(role='assistant', content=[MessageContent(text='')])),
                stop_reason='end_turn', usage=TokenUsage(), metrics=ConverseMetrics(latency_ms=latency_ms))
        choice = completion.choices[0]
        content = []
        reasoning = getattr(choice.message, 'reasoning', None) or getattr(choice.message, 'reasoning_content', None)
        if reasoning:
            content.append(MessageContent(reasoning_content=ReasoningContent(
                reasoning_text=ReasoningText(text=reasoning, signature=''), redacted_content=b'')))
        if choice.message.content:
            content.append(MessageContent(text=choice.message.content))
        for tc in (choice.message.tool_calls or []):
            args = tc.function.arguments
            if isinstance(args, str):
                try: args = json.loads(args)
                except json.JSONDecodeError: args = {"raw_input": args}
            content.append(MessageContent(tool_use=ToolUse(tool_use_id=tc.id, name=tc.function.name, input=args)))
        usage = completion.usage
        cache_read = 0
        if usage:
            details = getattr(usage, 'prompt_tokens_details', None)
            if details:
                cache_read = getattr(details, 'cached_tokens', 0) or 0
        return ConverseResponse(
            output=ConverseOutput(message=Message(role='assistant', content=content)),
            stop_reason=STOP_REASON_MAP.get(choice.finish_reason or '', 'end_turn'),
            usage=TokenUsage(input_tokens=usage.prompt_tokens if usage else 0,
                             output_tokens=usage.completion_tokens if usage else 0,
                             total_tokens=usage.total_tokens if usage else 0,
                             cache_read_input_tokens=cache_read),
            metrics=ConverseMetrics(latency_ms=latency_ms))

    def _consume_stream(self, stream, start) -> ConverseResponse:
        """Consume an OpenAI streaming response and assemble into ConverseResponse."""
        text_parts = []
        reasoning_parts = []
        tool_calls = {}  # index -> {id, name, arguments}
        finish_reason = None
        usage = None

        for chunk in stream:
            if chunk.usage:
                usage = chunk.usage
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            if chunk.choices[0].finish_reason:
                finish_reason = chunk.choices[0].finish_reason

            # Text content
            if delta.content:
                text_parts.append(delta.content)

            # Reasoning (provider-specific)
            reasoning = getattr(delta, 'reasoning', None) or getattr(delta, 'reasoning_content', None)
            if reasoning:
                reasoning_parts.append(reasoning)

            # Tool calls arrive incrementally by index
            for tc in (delta.tool_calls or []):
                idx = tc.index
                if idx not in tool_calls:
                    tool_calls[idx] = {'id': tc.id or '', 'name': '', 'arguments': ''}
                if tc.id:
                    tool_calls[idx]['id'] = tc.id
                if tc.function:
                    if tc.function.name:
                        tool_calls[idx]['name'] = tc.function.name
                    if tc.function.arguments:
                        tool_calls[idx]['arguments'] += tc.function.arguments

        # Build content list
        content = []
        if reasoning_parts:
            content.append(MessageContent(reasoning_content=ReasoningContent(
                reasoning_text=ReasoningText(text=''.join(reasoning_parts), signature=''), redacted_content=b'')))
        full_text = ''.join(text_parts)
        if full_text:
            content.append(MessageContent(text=full_text))
        for idx in sorted(tool_calls):
            tc = tool_calls[idx]
            args = tc['arguments']
            try:
                args = json.loads(args)
            except (json.JSONDecodeError, ValueError):
                args = {"raw_input": args}
            content.append(MessageContent(tool_use=ToolUse(tool_use_id=tc['id'], name=tc['name'], input=args)))

        cache_read = 0
        if usage:
            details = getattr(usage, 'prompt_tokens_details', None)
            if details:
                cache_read = getattr(details, 'cached_tokens', 0) or 0

        latency_ms = int((time.time() - start) * 1000)
        return ConverseResponse(
            output=ConverseOutput(message=Message(role='assistant', content=content)),
            stop_reason=STOP_REASON_MAP.get(finish_reason or '', 'end_turn'),
            usage=TokenUsage(
                input_tokens=getattr(usage, 'prompt_tokens', 0) if usage else 0,
                output_tokens=getattr(usage, 'completion_tokens', 0) if usage else 0,
                total_tokens=getattr(usage, 'total_tokens', 0) if usage else 0,
                cache_read_input_tokens=cache_read),
            metrics=ConverseMetrics(latency_ms=latency_ms))

    def with_thinking(self, tokens: int | str = 1024, enabled: bool = True):
        thinking_config = ThinkingConfig(
            type="enabled" if enabled else "disabled",
            budget_tokens=tokens
        )
        if self.additional_model_request_fields is None:
            self.additional_model_request_fields = AdditionalModelRequestFields()
        self.additional_model_request_fields.thinking = thinking_config
        return self

    def _get_response(self, messages=None):
        for callback in self.callbacks:
            try: callback.on_converse_start(self)
            except Exception as e: logger.warning(f"Callback error: {e}")
        params = self._build_params(messages)
        params['stream'] = True
        params['stream_options'] = {'include_usage': True}
        start = time.time()
        try:
            stream = self.openai_client.chat.completions.create(**params)
            response = self._consume_stream(stream, start)
        except Exception as error:
            for callback in self.callbacks:
                try:
                    if hasattr(callback, 'on_converse_error'):
                        callback.on_converse_error(self, error)
                except Exception as callback_error:
                    logger.warning(f"Callback error: {callback_error}")
            raise
        response.model_id = self.model_id
        for callback in self.callbacks:
            try: callback.on_converse_end(response)
            except Exception as e: logger.warning(f"Callback error: {e}")
        return response

    async def _aget_response(self, messages=None):
        for callback in self.callbacks:
            try: callback.on_converse_start(self)
            except Exception as e: logger.warning(f"Callback error: {e}")
        params = self._build_params(messages)
        params['stream'] = True
        params['stream_options'] = {'include_usage': True}
        start = time.time()
        try:
            stream = await self.async_openai_client.chat.completions.create(**params)
            # Async stream — collect chunks
            text_parts, reasoning_parts, tool_calls_map, finish_reason, usage = [], [], {}, None, None
            async for chunk in stream:
                if chunk.usage: usage = chunk.usage
                if not chunk.choices: continue
                delta = chunk.choices[0].delta
                if chunk.choices[0].finish_reason: finish_reason = chunk.choices[0].finish_reason
                if delta.content: text_parts.append(delta.content)
                reasoning = getattr(delta, 'reasoning', None) or getattr(delta, 'reasoning_content', None)
                if reasoning: reasoning_parts.append(reasoning)
                for tc in (delta.tool_calls or []):
                    idx = tc.index
                    if idx not in tool_calls_map:
                        tool_calls_map[idx] = {'id': tc.id or '', 'name': '', 'arguments': ''}
                    if tc.id: tool_calls_map[idx]['id'] = tc.id
                    if tc.function:
                        if tc.function.name: tool_calls_map[idx]['name'] = tc.function.name
                        if tc.function.arguments: tool_calls_map[idx]['arguments'] += tc.function.arguments

            content = []
            if reasoning_parts:
                content.append(MessageContent(reasoning_content=ReasoningContent(
                    reasoning_text=ReasoningText(text=''.join(reasoning_parts), signature=''), redacted_content=b'')))
            full_text = ''.join(text_parts)
            if full_text: content.append(MessageContent(text=full_text))
            for idx in sorted(tool_calls_map):
                tc = tool_calls_map[idx]
                args = tc['arguments']
                try: args = json.loads(args)
                except (json.JSONDecodeError, ValueError): args = {"raw_input": args}
                content.append(MessageContent(tool_use=ToolUse(tool_use_id=tc['id'], name=tc['name'], input=args)))

            cache_read = 0
            if usage:
                details = getattr(usage, 'prompt_tokens_details', None)
                if details: cache_read = getattr(details, 'cached_tokens', 0) or 0

            response = ConverseResponse(
                output=ConverseOutput(message=Message(role='assistant', content=content)),
                stop_reason=STOP_REASON_MAP.get(finish_reason or '', 'end_turn'),
                usage=TokenUsage(
                    input_tokens=getattr(usage, 'prompt_tokens', 0) if usage else 0,
                    output_tokens=getattr(usage, 'completion_tokens', 0) if usage else 0,
                    total_tokens=getattr(usage, 'total_tokens', 0) if usage else 0,
                    cache_read_input_tokens=cache_read),
                metrics=ConverseMetrics(latency_ms=int((time.time() - start) * 1000)))
        except Exception as error:
            for callback in self.callbacks:
                try:
                    if hasattr(callback, 'on_converse_error'):
                        callback.on_converse_error(self, error)
                except Exception as callback_error:
                    logger.warning(f"Callback error: {callback_error}")
            raise
        response.model_id = self.model_id
        for callback in self.callbacks:
            try: callback.on_converse_end(response)
            except Exception as e: logger.warning(f"Callback error: {e}")
        return response


@dataclass
class Mantle(_MantleTransport, Converse):
    api_key: Optional[str] = None


@dataclass
class MantleAgent(_MantleTransport, ConverseAgent):
    api_key: Optional[str] = None


@dataclass
class StructuredMantle(_MantleTransport, StructuredConverse):
    api_key: Optional[str] = None
