from __future__ import annotations

import json
import time
import threading
import urllib.error
import urllib.request
from typing import Any, Callable, Iterable, Iterator


class ProviderError(RuntimeError):
    def __init__(self, code: str, message: str, retryable: bool = False, status: int | None = None):
        super().__init__(message)
        self.code, self.retryable, self.status = code, retryable, status


def build_payload(model_id: str, messages: list[dict[str, Any]], preferences: dict[str, Any], model_key: str) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "model": model_id,
        "messages": messages,
        "stream": True,
        "temperature": preferences["temperature"],
        "top_p": preferences["top_p"],
        "max_tokens": preferences["max_output_tokens"],
    }
    if model_key == "k2.6":
        payload["reasoning"] = {"enabled": bool(preferences["reasoning_enabled"])}
    elif model_key == "k2.7-code":
        payload["reasoning"] = {"enabled": True}
        payload["chat_template_kwargs"] = {"clear_thinking": False}
    return payload


def parse_sse_lines(lines: Iterable[bytes]) -> Iterator[dict[str, str]]:
    for raw in lines:
        line = raw.decode("utf-8", errors="replace").strip()
        if not line or line.startswith(":") or not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if data == "[DONE]":
            yield {"type": "done", "text": ""}
            return
        try:
            obj = json.loads(data)
        except json.JSONDecodeError:
            continue
        if obj.get("error"):
            err = obj["error"]
            raise ProviderError("KIMU_PROVIDER_ERROR", str(err.get("message", "Provider error")), False)
        choices = obj.get("choices") or []
        if not choices:
            continue
        delta = choices[0].get("delta") or {}
        reasoning = delta.get("reasoning") or delta.get("reasoning_content")
        content = delta.get("content")
        if reasoning:
            yield {"type": "reasoning", "text": str(reasoning)}
        if content:
            yield {"type": "content", "text": str(content)}
        finish = choices[0].get("finish_reason")
        if finish:
            yield {"type": "finish", "text": str(finish)}


def fake_stream(model_key: str, user_text: str) -> Iterator[dict[str, str]]:
    reasoning = "Checking the selected model profile and preparing a deterministic local test response. "
    if model_key == "k2.7-code":
        reasoning += "K2.7 Code keeps thinking and preserved-reasoning controls locked on."
    else:
        reasoning += "K2.6 supports a user-selectable reasoning mode."
    for part in reasoning.split(" "):
        yield {"type": "reasoning", "text": part + " "}
        time.sleep(0.01)
    answer = f"KIMU chat is working in fake-provider mode with {model_key}. You said: {user_text}"
    for part in answer.split(" "):
        yield {"type": "content", "text": part + " "}
        time.sleep(0.01)
    yield {"type": "finish", "text": "stop"}


def stream_together(base_url: str, api_key: str, payload: dict[str, Any], timeout: int = 180, cancel_event: threading.Event | None = None, response_callback: Callable[[Any | None], None] | None = None) -> Iterator[dict[str, str]]:
    if not api_key:
        raise ProviderError("KIMU_PROVIDER_KEY_MISSING", "Together API key is not configured.", False)
    req = urllib.request.Request(
        base_url.rstrip("/") + "/chat/completions",
        data=json.dumps(payload).encode("utf-8"),
        method="POST",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
            "User-Agent": "KIMU/5.2.0-chat-foundation",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            if response_callback:
                response_callback(response)
            try:
                for item in parse_sse_lines(response):
                    if cancel_event and cancel_event.is_set():
                        return
                    yield item
            finally:
                if response_callback:
                    response_callback(None)
    except urllib.error.HTTPError as exc:
        body = exc.read(8192).decode("utf-8", errors="replace")
        message = "Together request failed."
        try:
            parsed = json.loads(body)
            message = str((parsed.get("error") or {}).get("message") or message)
        except json.JSONDecodeError:
            pass
        retryable = exc.code == 429 or 500 <= exc.code < 600
        code = "KIMU_PROVIDER_RATE_LIMIT" if exc.code == 429 else "KIMU_PROVIDER_HTTP_ERROR"
        raise ProviderError(code, message, retryable, exc.code) from exc
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        if cancel_event and cancel_event.is_set():
            return
        raise ProviderError("KIMU_PROVIDER_UNREACHABLE", "Together AI could not be reached.", True) from exc
    finally:
        if response_callback:
            response_callback(None)


def complete_together(
    base_url: str,
    api_key: str,
    payload: dict[str, Any],
    timeout: int = 180,
) -> dict[str, Any]:
    """Run one non-streaming OpenAI-compatible Together chat completion.

    This is used by bounded tool loops where the complete assistant message,
    including tool_calls, must be validated before any local side effect.
    """
    if not api_key:
        raise ProviderError("KIMU_PROVIDER_KEY_MISSING", "Together API key is not configured.", False)
    body = dict(payload)
    body["stream"] = False
    req = urllib.request.Request(
        base_url.rstrip("/") + "/chat/completions",
        data=json.dumps(body).encode("utf-8"),
        method="POST",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
            "User-Agent": "KIMU/5.2.0-tool-runtime",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            parsed = json.loads(response.read(16 * 1024 * 1024).decode("utf-8"))
    except urllib.error.HTTPError as exc:
        raw = exc.read(8192).decode("utf-8", errors="replace")
        message = "Together request failed."
        try:
            message = str((json.loads(raw).get("error") or {}).get("message") or message)
        except json.JSONDecodeError:
            pass
        retryable = exc.code == 429 or 500 <= exc.code < 600
        raise ProviderError("KIMU_PROVIDER_HTTP_ERROR", message, retryable, exc.code) from exc
    except (urllib.error.URLError, TimeoutError, OSError, json.JSONDecodeError) as exc:
        raise ProviderError("KIMU_PROVIDER_UNREACHABLE", "Together AI could not be reached or returned invalid JSON.", True) from exc
    choices = parsed.get("choices") or []
    if not choices or not isinstance(choices[0].get("message"), dict):
        raise ProviderError("KIMU_PROVIDER_RESPONSE_INVALID", "Together returned no assistant message.", False)
    message = choices[0]["message"]
    usage = parsed.get("usage") if isinstance(parsed.get("usage"), dict) else {}
    return {"message": message, "usage": usage, "finish_reason": choices[0].get("finish_reason")}


def fake_tool_completion(messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> dict[str, Any]:
    """Deterministic no-cost tool-loop response used by verification."""
    last = messages[-1] if messages else {"role": "user", "content": ""}
    if last.get("role") == "tool":
        return {
            "message": {"role": "assistant", "content": "The requested workspace operation completed and its audited result is available."},
            "usage": {"prompt_tokens": 24, "completion_tokens": 18, "total_tokens": 42},
            "finish_reason": "stop",
        }
    prompt = str(last.get("content") or "").lower()
    selected = "workspace.list"
    arguments: dict[str, Any] = {"path": ".", "recursive": True, "limit": 100}
    if "read " in prompt or "show " in prompt:
        selected, arguments = "workspace.read", {"path": "README.md", "max_chars": 50000}
    elif "create" in prompt or "write" in prompt:
        selected, arguments = "workspace.write", {"path": "notes/tool-agent-result.md", "content": "Created by the bounded KIMU tool agent.\n", "expected_sha256": None}
    elif "search" in prompt or "find" in prompt:
        selected, arguments = "workspace.search", {"query": "KIMU", "path": ".", "regex": False, "case_sensitive": False, "limit": 100}
    available = {item.get("function", {}).get("name") for item in tools}
    if selected not in available:
        return {"message": {"role": "assistant", "content": "No compatible tool was available."}, "usage": {"total_tokens": 20}, "finish_reason": "stop"}
    call_id = "call_fake_" + str(int(time.time() * 1000))
    return {
        "message": {
            "role": "assistant",
            "content": "",
            "tool_calls": [{"id": call_id, "type": "function", "function": {"name": selected, "arguments": json.dumps(arguments)}}],
        },
        "usage": {"prompt_tokens": 32, "completion_tokens": 20, "total_tokens": 52},
        "finish_reason": "tool_calls",
    }