import json
import socket
from lib import gadget
import urllib.request
import urllib.error

class _LLMFail(Exception):
    """Kegagalan network dari satu kandidat chain.

    kind:
      - 'auth'        : 401/403 → rotasi key (coba key berikutnya di provider yang sama)
      - 'unreachable' : koneksi gagal / timeout / 5xx / 429 / 404 → pindah model/provider
    """
    def __init__(self, kind, message):
        super().__init__(message)
        self.kind = kind
        self.message = message


class LLMClient:
    class Message:
        def __init__(self, msg):
            raw_content     = msg.get('content', '')      # Ambil konten mentah
            self.content    = gadget.strip_thinking(raw_content) if isinstance(raw_content, str) else raw_content # Auto-strip <thinking> dari content
            self.tool_calls = msg.get('tool_calls', None) # Ambil tool calls
            self.warning    = None

    def __init__(self, base_url, model, api_key, timeout=600):
        self.base_url = (base_url or "").rstrip('/')
        self.model = model
        self.api_key = api_key
        self.timeout = timeout
        self.cancel_requested = False

        # Chain kandidat per tipe ('llm', 'imagevision', ...)
        self._chains: dict[str, list[dict]] = {}
        self._chain: list[dict] | None = None
        self._chain_key: str | None = None
        self._chain_index = 0

    # ------------------------------------------------------------ config ---

    def set_chains(self, chains):
        """Set chain kandidat per tipe. Memperbarui base_url/model/api_key dari
        kandidat pertama chain 'llm' (atau 'imagevision' sebagai fallback)."""
        self._chains = chains or {}
        self._chain = None
        self._chain_key = None
        self._chain_index = 0
        llm = self._chains.get("llm") or []
        vision = self._chains.get("imagevision") or []
        if llm:
            self._apply_candidate(llm[0])
        elif vision:
            self._apply_candidate(vision[0])
        else:
            self.base_url = ""
            self.model = ""
            self.api_key = ""

    def has_chain(self, type_name="llm"):
        return bool(self._chains.get(type_name))

    def _apply_candidate(self, cand):
        self.base_url = (cand.get("base_url") or "").rstrip('/')
        self.model = cand.get("model", "")
        self.api_key = cand.get("api_key", "")

    # ----------------------------------------------------- chain selection ---

    def _pick_chain_type(self, messages):
        """Kalau ada konten gambar (image_url), pakai chain imagevision."""
        for m in messages or []:
            content = m.get("content")
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "image_url":
                        return "imagevision"
        return "llm"

    def chat(self, messages, tools=None, on_stream_chunk=None, disable_reasoning=False, tool_reminder=None):
        if not self._chains:
            return self._chat_once(
                self.base_url, self.model, self.api_key, messages,
                tools=tools, on_stream_chunk=on_stream_chunk,
                disable_reasoning=disable_reasoning, tool_reminder=tool_reminder,
            )

        chain_type = self._pick_chain_type(messages)
        if chain_type != self._chain_key:
            self._chain_key = chain_type
            self._chain_index = 0
        chain = self._chains.get(chain_type) or []
        self._chain = chain

        if not chain:
            # Vision diminta tapi tidak ada chain imagevision → pakai llm.
            if chain_type == "imagevision" and self._chains.get("llm"):
                chain = self._chains["llm"]
                self._chain = chain
                self._chain_key = "llm"
            else:
                return self.Message({
                    'content': "Error: Tidak ada model yang dikonfigurasi untuk permintaan ini.",
                    'tool_calls': None,
                })

        return self._chat_chain(
            messages, tools=tools, on_stream_chunk=on_stream_chunk,
            disable_reasoning=disable_reasoning, tool_reminder=tool_reminder,
        )

    # ---------------------------------------------------------- chain loop ---

    @staticmethod
    def _same_target(a, b):
        return a.get("api_id") == b.get("api_id") and a.get("model_id") == b.get("model_id")

    @staticmethod
    def _next_target(chain, idx):
        """Lompat ke kandidat model/provider berikutnya (skip sisa key)."""
        api = chain[idx].get("api_id")
        model = chain[idx].get("model_id")
        j = idx + 1
        while j < len(chain) and chain[j].get("api_id") == api and chain[j].get("model_id") == model:
            j += 1
        return j

    def _chat_chain(self, messages, tools=None, on_stream_chunk=None,
                    disable_reasoning=False, tool_reminder=None):
        chain = self._chain or []
        if not chain:
            return self.Message({'content': "Error: Tidak ada model yang dikonfigurasi.", 'tool_calls': None})

        idx = min(max(self._chain_index, 0), len(chain) - 1)
        last = None

        while idx < len(chain):
            cand = chain[idx]
            self._apply_candidate(cand)
            try:
                result = self._chat_once(
                    cand["base_url"], cand["model"], cand["api_key"], messages,
                    tools=tools, on_stream_chunk=on_stream_chunk,
                    disable_reasoning=disable_reasoning, tool_reminder=tool_reminder,
                )
                self._chain_index = idx
                return result
            except _LLMFail as e:
                last = e
                if e.kind == "auth" and idx + 1 < len(chain) and self._same_target(chain[idx + 1], cand):
                    idx += 1  # rotasi key
                else:
                    idx = self._next_target(chain, idx)  # pindah model/provider

        # Image vision gagal total → fallback ke model llm umum.
        if self._chain_key == "imagevision":
            fb = self._chains.get("llm") or []
            if fb:
                self._chain = fb
                self._chain_key = "llm"
                self._chain_index = 0
                result = self._chat_chain(
                    messages, tools=tools, on_stream_chunk=on_stream_chunk,
                    disable_reasoning=disable_reasoning, tool_reminder=tool_reminder,
                )
                result.warning = (f"{result.warning} " if result.warning else "") + \
                    "Image vision model gagal, fallback ke LLM umum."
                return result

        msg = "Error: Semua opsi model gagal."
        if last:
            msg += f"\n{last.message}"
        return self.Message({'content': msg, 'tool_calls': None})

    # ------------------------------------------------------- single attempt ---

    def _chat_once(self, base_url, model, api_key, messages, tools=None,
                   on_stream_chunk=None, disable_reasoning=False, tool_reminder=None):
        url = f"{base_url.rstrip('/')}/chat/completions"
        payload = {
            "model": model,
            "messages": messages,
            "stream": True  # Enable streaming
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"
            if tool_reminder:
                # Inject tool reminder sebagai system message sebelum user message terakhir
                for i in range(len(messages) - 1, -1, -1):
                    if messages[i].get("role") == "user":
                        payload["messages"] = messages[:i] + [
                            {"role": "system", "content": tool_reminder}
                        ] + messages[i:]
                        break

        # Hanya kirim parameter reasoning jika diminta eksplisit
        # Beberapa model/provider justru error jika parameter ini ada tapi tidak didukung
        if disable_reasoning:
            payload["reasoning"] = {"enabled": False}

        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(url, data=data, method='POST')
        req.add_header('Content-Type', 'application/json')
        if api_key:
            req.add_header('Authorization', f'Bearer {api_key}')

        # Variabel untuk mengumpulkan hasil
        full_content = ""
        full_tool_calls = []
        reasoning_content = ""
        full_multimodal_content = []
        stream_started = False

        try:
            self.cancel_requested = False
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                # Streaming: baca line by line
                for line in resp:
                    if self.cancel_requested:
                        # Stream cancellation
                        full_content += "\n\n[Stream cancelled by user]"
                        break

                    line = line.decode('utf-8').strip()
                    if not line or not line.startswith('data: '):
                        continue

                    data_str = line[6:]  # Hapus "data: " prefix
                    if data_str == '[DONE]':
                        break

                    try:
                        chunk = json.loads(data_str)
                    except json.JSONDecodeError:
                        continue

                    # Parse delta dari chunk
                    delta = chunk.get('choices', [{}])[0].get('delta', {})
                    finish_reason = chunk.get('choices', [{}])[0].get('finish_reason', None)

                    # Stream reasoning content jika ada
                    if 'reasoning_content' in delta:
                        reasoning_content += delta['reasoning_content']

                    # Stream tool_calls jika ada
                    if 'tool_calls' in delta:
                        tool_calls = delta['tool_calls']
                        for tc in tool_calls:
                            idx = tc.get('index', 0)
                            # Pastikan list cukup panjang
                            while len(full_tool_calls) <= idx:
                                full_tool_calls.append({
                                    "id": "",
                                    "type": "function",
                                    "function": {"name": "", "arguments": ""}
                                })

                            # Update ID
                            if 'id' in tc and tc['id']:
                                full_tool_calls[idx]['id'] = tc['id']

                            # Update function name
                            if 'function' in tc and 'name' in tc['function']:
                                full_tool_calls[idx]['function']['name'] += tc['function']['name']

                            # Update arguments
                            if 'function' in tc and 'arguments' in tc['function']:
                                full_tool_calls[idx]['function']['arguments'] += tc['function']['arguments']

                    # Stream content (text response atau multimodal blocks)
                    if 'content' in delta:
                        chunk = delta['content']
                        if isinstance(chunk, str):
                            full_content += chunk or ""
                            if chunk:
                                stream_started = True
                            if on_stream_chunk and chunk:
                                on_stream_chunk(chunk)
                        elif isinstance(chunk, list):
                            full_multimodal_content.extend(chunk)
                            for block in chunk:
                                if isinstance(block, dict) and block.get('type') == 'text':
                                    t = block.get('text', '')
                                    if t:
                                        full_content += t
                                        stream_started = True
                                        if on_stream_chunk:
                                            on_stream_chunk(t)

            # Build final response
            if full_multimodal_content:
                message = {'content': full_multimodal_content}
            else:
                message = {'content': full_content}

            if full_tool_calls:
                # Filter tool_calls yang valid (ada name dan arguments)
                valid_tool_calls = []
                for tc in full_tool_calls:
                    name = tc.get('function', {}).get('name')
                    args_str = tc.get('function', {}).get('arguments')
                    tc_id = tc.get('id')

                    # Pastikan name dan arguments ada
                    if name and args_str and args_str.strip():
                        # Generate ID jika kosong
                        if not tc_id:
                            tc_id = f"call_{len(valid_tool_calls)}"
                        tc['id'] = tc_id

                        try:
                            # Validate dan re-encode JSON untuk format yang konsisten
                            parsed_args = json.loads(args_str)
                            tc['function']['arguments'] = json.dumps(parsed_args, ensure_ascii=False)
                            valid_tool_calls.append(tc)
                        except json.JSONDecodeError:
                            # Invalid JSON, coba raw string tapi hanya jika tidak kosong
                            tc['function']['arguments'] = args_str
                            valid_tool_calls.append(tc)

                if valid_tool_calls:
                    message['tool_calls'] = valid_tool_calls

            response = {'choices': [{'message': message}]}
        except urllib.error.HTTPError as e:
            body_text = ""
            try:
                body_text = e.read().decode('utf-8', errors='replace')
            except Exception:
                pass
            if tools and e.code == 404:
                try:
                    body = json.loads(body_text) if body_text else {}
                    if 'tool use' in body.get('error', {}).get('message', '').lower():
                        result = self._chat_once(
                            base_url, model, api_key, messages, tools=None,
                            on_stream_chunk=on_stream_chunk,
                            disable_reasoning=disable_reasoning,
                        )
                        result.warning = "Tool calling not supported by this model. Running in chat-only mode."
                        return result
                except Exception:
                    pass
            detail = f" - {body_text[:500]}" if body_text else ""
            if e.code in (401, 403):
                raise _LLMFail("auth", f"HTTP {e.code} {e.reason}{detail}")
            raise _LLMFail("unreachable", f"HTTP {e.code} {e.reason}{detail}")
        except urllib.error.URLError as e:
            if stream_started:
                return self.Message({'content': f"Error: {e.reason}", 'tool_calls': None})
            raise _LLMFail("unreachable", f"Connection error: {e.reason}")
        except (TimeoutError, socket.timeout) as e:
            if stream_started:
                return self.Message({'content': f"Error: Timeout", 'tool_calls': None})
            raise _LLMFail("unreachable", f"Timeout")
        except Exception as e:
            if stream_started:
                return self.Message({'content': f"Error: {str(e)}", 'tool_calls': None})
            raise _LLMFail("unreachable", f"Error: {str(e)}")

        if 'choices' not in response:
            raw_preview = json.dumps(response)[:500]
            raise _LLMFail("unreachable", (
                f"Unexpected response — 'choices' key missing.\n"
                f"  URL   : {url}\n"
                f"  Model : {model}\n"
                f"  Response: {raw_preview}"
            ))
        if not response['choices']:
            raw_preview = json.dumps(response)[:500]
            raise _LLMFail("unreachable", (
                f"'choices' is empty in the response.\n"
                f"  URL   : {url}\n"
                f"  Model : {model}\n"
                f"  Response: {raw_preview}"
            ))
        if 'message' not in response['choices'][0]:
            raw_preview = json.dumps(response['choices'][0])[:500]
            raise _LLMFail("unreachable", (
                f"'message' key missing in first choice.\n"
                f"  URL   : {url}\n"
                f"  Model : {model}\n"
                f"  Choice : {raw_preview}"
            ))

        message = response['choices'][0]['message']

        # Handle reasoning_content field dari OpenRouter/models yang support thinking
        # Pindahkan ke content jangan sampai keluar
        message.pop('reasoning_content', None)
        message.pop('reasoning', None)

        return self.Message(message)