from __future__ import annotations

import json
import urllib.error
import urllib.request
from typing import Any


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


def execute_tci(
    base_url: str,
    api_key: str,
    code: str,
    session_id: str | None = None,
    files: list[dict[str, Any]] | None = None,
    timeout: int = 180,
) -> dict[str, Any]:
    if not api_key:
        raise TCIError("KIMU_PROVIDER_KEY_MISSING", "Together API key is not configured.")
    payload: dict[str, Any] = {"code": code, "language": "python"}
    if session_id:
        payload["session_id"] = session_id
    if files:
        payload["files"] = files
    request = urllib.request.Request(
        base_url.rstrip("/") + "/tci/execute",
        data=json.dumps(payload).encode("utf-8"),
        method="POST",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
            "User-Agent": "KIMU/5.2.0-increment-39",
        },
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            parsed = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read(8192).decode("utf-8", errors="replace")
        message = "Together Code Interpreter request failed."
        try:
            obj = json.loads(body)
            errors = obj.get("errors")
            if isinstance(errors, list) and errors:
                message = str(errors[0])
            message = str((obj.get("error") or {}).get("message") or message)
        except json.JSONDecodeError:
            pass
        raise TCIError(
            "KIMU_TCI_RATE_LIMIT" if exc.code == 429 else "KIMU_TCI_HTTP_ERROR",
            message,
            exc.code == 429 or 500 <= exc.code < 600,
            exc.code,
        ) from exc
    except (urllib.error.URLError, TimeoutError) as exc:
        raise TCIError("KIMU_TCI_UNREACHABLE", "Together Code Interpreter could not be reached.", True) from exc
    if parsed.get("errors"):
        raise TCIError("KIMU_TCI_EXECUTION_ERROR", str(parsed["errors"]), False)
    data = parsed.get("data") or {}
    return {
        "status": data.get("status") or "completed",
        "session_id": data.get("session_id"),
        "outputs": data.get("outputs") or [],
    }


def fake_execute(code: str, session_id: str | None = None) -> dict[str, Any]:
    return {
        "status": "completed",
        "session_id": session_id or "ses_fake_kimu",
        "outputs": [
            {
                "type": "stdout",
                "data": "KIMU Code Interpreter fake-provider mode completed.\n"
                + f"Received {len(code)} characters of Python code.\n",
            }
        ],
    }