"""
OpenCode ↔ Anthropic Gateway  —  latency-optimized edition
"""

import os
import hashlib
import base64
import uuid
import asyncio
import logging
from contextlib import asynccontextmanager

try:
    import orjson as _json_lib
    def _dumps(obj: object) -> str:
        return _json_lib.dumps(obj).decode()
    def _loads(s):
        return _json_lib.loads(s)
except ImportError:
    import json as _json_lib        # type: ignore[assignment]
    def _dumps(obj: object) -> str:
        return _json_lib.dumps(obj, separators=(",", ":"))
    def _loads(s):
        return _json_lib.loads(s)

import json
import httpx

from fastapi import FastAPI, Request, Header
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("OpenCodeGateway")

OPENCODE_API_KEY = os.getenv("OPENCODE_API_KEY", "").strip()
UPSTREAM_URL     = "https://opencode.ai/zen/v1/chat/completions"
MAX_RETRIES      = 2
THINKING_MODE    = os.getenv("THINKING_MODE", "enabled").strip().lower()

_OPENCODE_VERSION = os.getenv("OPENCODE_VERSION", "1.15.4")
_USER_ID    = os.getenv("OPENCODE_USER_ID",    uuid.uuid4().hex)
_PROJECT_ID = os.getenv("OPENCODE_PROJECT_ID", uuid.uuid4().hex)

_STATIC_UPSTREAM_HEADERS = {
    "Content-Type":      "application/json",
    "Accept-Encoding":   "identity",
    "Accept":            "application/json",
    "Origin":            "https://opencode.ai",
    "Referer":           "https://opencode.ai/",
    "User-Agent":        f"opencode/{_OPENCODE_VERSION}",
    "x-opencode-client": "cli",
}

def _make_upstream_headers(api_key: str) -> dict:
    return {
        **_STATIC_UPSTREAM_HEADERS,
        "x-api-key":           api_key,
        "x-opencode-request":  _USER_ID,
        "x-opencode-project":  _PROJECT_ID,
        "x-opencode-session":  uuid.uuid4().hex,
    }

STREAM_RESPONSE_HEADERS = {
    "Cache-Control":     "no-cache",
    "Connection":        "keep-alive",
    "X-Accel-Buffering": "no",
}

# =========================================================
# MODELS
# =========================================================
# النماذج المجانية الحالية في OpenCode Zen (مُحدَّثة 2026-06):
#   deepseek-v4-flash  → DeepSeek V4 Flash Free  (الأقوى والأسرع)
#   mimo-v2.5          → MiMo-V2.5 Free          (سريع جداً)
#   north-mini-code    → North Mini Code Free     (مخصص للكود)
#   big-pickle         → Big Pickle Free          (مجاني دائم)
#
# الأسماء القديمة (nemotron-3-super-free, deepseek-v4-flash-free …)
# لم تعد موجودة — الأسماء الصحيحة الآن بدون لاحقة "-free".

MODEL_MAP = {
    # Opus → أقوى نموذج مجاني متاح
    "claude-opus-4-6":           "deepseek-v4-flash",
    "claude-opus-4-5":           "deepseek-v4-flash",
    # Sonnet → DeepSeek V4 Flash (أفضل توازن جودة/سرعة)
    "claude-sonnet-4-6":         "deepseek-v4-flash",
    "claude-sonnet-4-5":         "deepseek-v4-flash",
    # Haiku → MiMo-V2.5 (أسرع للمهام الخفيفة)
    "claude-haiku-4-6":          "mimo-v2.5",
    "claude-haiku-4-5":          "mimo-v2.5",
    "claude-haiku-4-5-20251001": "mimo-v2.5",
}
DEFAULT_MODEL = "deepseek-v4-flash"

_MODELS_CACHE = {
    "object": "list",
    "data": [
        {"id": m, "object": "model", "created": 1710000000, "owned_by": "gateway"}
        for m in MODEL_MAP
    ],
}

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.client = httpx.AsyncClient(
        http2=True,
        timeout=httpx.Timeout(connect=15, read=300, write=60, pool=30),
        limits=httpx.Limits(max_connections=100, max_keepalive_connections=30),
        follow_redirects=True,
    )
    logger.info("Gateway started — THINKING_MODE=%s  json_backend=%s",
                THINKING_MODE, _json_lib.__name__)
    yield
    await app.state.client.aclose()


app = FastAPI(lifespan=lifespan)

@app.get("/")
@app.head("/")
async def root():
    return {"status": "ok", "service": "OpenCode Gateway"}

@app.get("/health")
async def health():
    return {"status": "healthy"}

@app.get("/v1/models")
async def get_models():
    return _MODELS_CACHE

def _fix_schema(schema):
    if not isinstance(schema, dict):
        return {"type": "object", "properties": {}}
    if schema.get("type") not in (
        "object", "array", "string", "number", "boolean", "integer"
    ):
        schema = dict(schema)
        schema["type"] = "object"
        schema.setdefault("properties", {})
    return schema

def _synthetic_signature(text: str) -> str:
    digest = hashlib.sha256(text.encode()).digest()
    return "synthetic_" + base64.urlsafe_b64encode(digest).decode().rstrip("=")

def normalize_messages(messages: list) -> list:
    output = []
    for msg in messages:
        role    = msg.get("role", "user")
        content = msg.get("content", "")
        reasoning_content = msg.get("reasoning_content")

        if isinstance(content, str):
            result = {"role": role, "content": content}
            if reasoning_content is not None:
                result["reasoning_content"] = reasoning_content
            output.append(result)
            continue

        if isinstance(content, list):
            text_parts    = []
            tool_calls    = []
            tool_results  = []
            thinking_parts = []

            for item in content:
                if not isinstance(item, dict):
                    text_parts.append(str(item))
                    continue
                t = item.get("type", "")
                if t == "text":
                    txt = item.get("text", "")
                    if txt: text_parts.append(txt)
                elif t == "thinking":
                    think_txt = item.get("thinking", "")
                    if think_txt: thinking_parts.append(think_txt)
                elif t == "tool_use":
                    tool_calls.append({
                        "id":   item.get("id", f"call_{uuid.uuid4().hex[:8]}"),
                        "type": "function",
                        "function": {
                            "name":      item.get("name", ""),
                            "arguments": json.dumps(item.get("input", {})),
                        },
                    })
                elif t == "tool_result":
                    rc = item.get("content", "")
                    if isinstance(rc, list):
                        rc = "\n".join(
                            c.get("text", "") if isinstance(c, dict) else str(c)
                            for c in rc
                        )
                    tool_results.append({
                        "role":         "tool",
                        "tool_call_id": item.get("tool_use_id", ""),
                        "content":      str(rc),
                    })

            final_reasoning = (
                reasoning_content if reasoning_content is not None
                else ("\n".join(thinking_parts) if thinking_parts else None)
            )

            if tool_calls:
                assistant_msg = {
                    "role":       "assistant",
                    "tool_calls": tool_calls,
                    "content":    "\n".join(text_parts) if text_parts else None,
                }
                if final_reasoning is not None:
                    assistant_msg["reasoning_content"] = final_reasoning
                output.append(assistant_msg)
            elif tool_results:
                output.extend(tool_results)
            else:
                result = {"role": role, "content": "\n".join(text_parts)}
                if final_reasoning is not None:
                    result["reasoning_content"] = final_reasoning
                output.append(result)
            continue

        result = {"role": role, "content": str(content)}
        if reasoning_content is not None:
            result["reasoning_content"] = reasoning_content
        output.append(result)

    return output

def convert_tools_to_openai(tools: list) -> list:
    if not tools:
        return []
    return [
        {
            "type": "function",
            "function": {
                "name":        t.get("name", ""),
                "description": t.get("description", ""),
                "parameters":  _fix_schema(t.get("input_schema", {})),
            },
        }
        for t in tools
    ]

def map_tool_choice(raw, has_tools: bool, thinking_on: bool):
    if not raw or not has_tools: return None
    if isinstance(raw, dict):
        kind = raw.get("type", "auto")
        if thinking_on and kind in ("any", "tool"):
            logger.warning("tool_choice '%s' incompatible with extended thinking → auto", kind)
            return "auto"
        if kind == "any":   return "required"
        if kind == "tool":  return {"type": "function", "function": {"name": raw.get("name", "")}}
        return kind
    if thinking_on and raw in ("any", "required"):
        return "auto"
    return raw

def build_payload(body, messages, tools, tool_choice, model, stream):
    thinking_on = THINKING_MODE == "enabled"
    payload = {
        "model":      model,
        "messages":   messages,
        "max_tokens": min(body.get("max_tokens", 32000), 64000),
        "stream":     stream,
        "thinking":   {"type": "enabled"} if thinking_on else {"type": "disabled"},
    }
    if not thinking_on:
        payload["temperature"] = body.get("temperature", 0.7)
    elif "temperature" in body:
        payload["temperature"] = body["temperature"]
    if tools:
        payload["tools"] = tools
        if tool_choice is not None:
            payload["tool_choice"] = tool_choice
    return payload

_FINISH_MAP = {"tool_calls": "tool_use", "length": "max_tokens"}

def convert_to_anthropic(parsed: dict, requested_model: str) -> dict:
    content_blocks = []
    stop_reason    = "end_turn"
    try:
        choice  = (parsed.get("choices") or [{}])[0]
        message = choice.get("message", {})
        stop_reason = _FINISH_MAP.get(choice.get("finish_reason"), "end_turn")
        reasoning = message.get("reasoning_content")
        if reasoning:
            content_blocks.append({
                "type": "thinking", "thinking": reasoning,
                "signature": _synthetic_signature(reasoning),
            })
        text = message.get("content") or ""
        if text:
            content_blocks.append({"type": "text", "text": text})
        for tc in (message.get("tool_calls") or []):
            fn = tc.get("function", {})
            try:    parsed_args = _loads(fn.get("arguments", "{}"))
            except: parsed_args = {"raw": fn.get("arguments", "")}
            content_blocks.append({
                "type":  "tool_use",
                "id":    tc.get("id", f"toolu_{uuid.uuid4().hex[:8]}"),
                "name":  fn.get("name", ""),
                "input": parsed_args,
            })
    except Exception as exc:
        logger.warning("convert_to_anthropic error: %s", exc)

    if not content_blocks:
        content_blocks.append({"type": "text", "text": ""})

    usage = parsed.get("usage", {})
    return {
        "id": f"msg_{uuid.uuid4().hex}", "type": "message",
        "role": "assistant", "model": requested_model,
        "content": content_blocks, "stop_reason": stop_reason,
        "stop_sequence": None,
        "usage": {
            "input_tokens":  usage.get("prompt_tokens", 0),
            "output_tokens": usage.get("completion_tokens", 0),
        },
    }

async def _iter_sse_lines(response):
    buf = b""
    async for chunk in response.aiter_bytes():
        if not chunk: continue
        buf += chunk
        while True:
            nl = buf.find(b"\n")
            if nl == -1: break
            line = buf[:nl].rstrip(b"\r").decode("utf-8", errors="replace")
            buf  = buf[nl + 1:]
            yield line

_MSG_START = "event: message_start\ndata: "
_BLK_START = "event: content_block_start\ndata: "
_BLK_DELTA = "event: content_block_delta\ndata: "
_BLK_STOP  = "event: content_block_stop\ndata: "
_MSG_DELTA = "event: message_delta\ndata: "
_MSG_STOP  = "event: message_stop\ndata: "

async def convert_stream_to_anthropic(response, requested_model: str):
    msg_id        = f"msg_{uuid.uuid4().hex}"
    input_tokens  = 0
    output_tokens = 0

    yield (f"{_MSG_START}" + _dumps({
        "type": "message_start",
        "message": {
            "id": msg_id, "type": "message", "role": "assistant",
            "model": requested_model, "content": [],
            "stop_reason": None, "stop_sequence": None,
            "usage": {"input_tokens": 0, "output_tokens": 0},
        },
    }) + "\n\n")
    yield 'event: ping\ndata: {"type":"ping"}\n\n'

    tool_state      = {}
    tool_block_idx  = {}
    next_idx        = 0
    reasoning_block_idx  = None
    reasoning_text_acc   = ""
    reasoning_closed     = False
    text_block_idx = None
    stop_reason    = "end_turn"

    try:
        async for line in _iter_sse_lines(response):
            if not line or not line.startswith("data:"): continue
            data = line[5:].strip()
            if data == "[DONE]": break
            try: chunk = _loads(data)
            except: continue

            if chunk.get("usage"):
                u             = chunk["usage"]
                input_tokens  = u.get("prompt_tokens",     input_tokens)
                output_tokens = u.get("completion_tokens", output_tokens)

            choices = chunk.get("choices")
            if not choices: continue
            choice        = choices[0]
            delta         = choice.get("delta", {})
            finish_reason = choice.get("finish_reason")
            if finish_reason == "tool_calls":  stop_reason = "tool_use"
            elif finish_reason == "length":    stop_reason = "max_tokens"

            reasoning_delta = delta.get("reasoning_content")
            if reasoning_delta:
                reasoning_text_acc += reasoning_delta
                if reasoning_block_idx is None:
                    reasoning_block_idx = next_idx; next_idx += 1
                    yield (_BLK_START + _dumps({"type": "content_block_start",
                        "index": reasoning_block_idx,
                        "content_block": {"type": "thinking", "thinking": ""}}) + "\n\n")
                yield (_BLK_DELTA + _dumps({"type": "content_block_delta",
                    "index": reasoning_block_idx,
                    "delta": {"type": "thinking_delta", "thinking": reasoning_delta}}) + "\n\n")

            text_delta = delta.get("content")
            if text_delta:
                if reasoning_block_idx is not None and not reasoning_closed:
                    sig = _synthetic_signature(reasoning_text_acc)
                    yield (_BLK_DELTA + _dumps({"type": "content_block_delta",
                        "index": reasoning_block_idx,
                        "delta": {"type": "signature_delta", "signature": sig}}) + "\n\n")
                    yield (_BLK_STOP + _dumps({"type": "content_block_stop",
                        "index": reasoning_block_idx}) + "\n\n")
                    reasoning_closed = True
                if text_block_idx is None:
                    text_block_idx = next_idx; next_idx += 1
                    yield (_BLK_START + _dumps({"type": "content_block_start",
                        "index": text_block_idx,
                        "content_block": {"type": "text", "text": ""}}) + "\n\n")
                yield (_BLK_DELTA + _dumps({"type": "content_block_delta",
                    "index": text_block_idx,
                    "delta": {"type": "text_delta", "text": text_delta}}) + "\n\n")

            tool_call_deltas = delta.get("tool_calls")
            if tool_call_deltas:
                if reasoning_block_idx is not None and not reasoning_closed:
                    sig = _synthetic_signature(reasoning_text_acc)
                    yield (_BLK_DELTA + _dumps({"type": "content_block_delta",
                        "index": reasoning_block_idx,
                        "delta": {"type": "signature_delta", "signature": sig}}) + "\n\n")
                    yield (_BLK_STOP + _dumps({"type": "content_block_stop",
                        "index": reasoning_block_idx}) + "\n\n")
                    reasoning_closed = True
                for tc_delta in tool_call_deltas:
                    idx = tc_delta.get("index", 0)
                    if idx not in tool_state:
                        fn_init = tc_delta.get("function", {})
                        tool_state[idx] = {
                            "id":   tc_delta.get("id", f"toolu_{uuid.uuid4().hex[:8]}"),
                            "name": fn_init.get("name", ""), "args": "",
                        }
                        blk = next_idx; next_idx += 1; tool_block_idx[idx] = blk
                        yield (_BLK_START + _dumps({"type": "content_block_start", "index": blk,
                            "content_block": {"type": "tool_use", "id": tool_state[idx]["id"],
                            "name": tool_state[idx]["name"], "input": {}}}) + "\n\n")
                    fn = tc_delta.get("function", {})
                    args = fn.get("arguments", "")
                    if args:
                        tool_state[idx]["args"] += args
                        yield (_BLK_DELTA + _dumps({"type": "content_block_delta",
                            "index": tool_block_idx[idx],
                            "delta": {"type": "input_json_delta", "partial_json": args}}) + "\n\n")

    except Exception as exc:
        logger.warning("stream error: %s", exc)

    if reasoning_block_idx is not None and not reasoning_closed:
        sig = _synthetic_signature(reasoning_text_acc)
        yield (_BLK_DELTA + _dumps({"type": "content_block_delta", "index": reasoning_block_idx,
            "delta": {"type": "signature_delta", "signature": sig}}) + "\n\n")
        yield (_BLK_STOP + _dumps({"type": "content_block_stop",
            "index": reasoning_block_idx}) + "\n\n")

    if text_block_idx is not None:
        yield (_BLK_STOP + _dumps({"type": "content_block_stop",
            "index": text_block_idx}) + "\n\n")

    for blk in tool_block_idx.values():
        yield (_BLK_STOP + _dumps({"type": "content_block_stop", "index": blk}) + "\n\n")

    yield (_MSG_DELTA + _dumps({"type": "message_delta",
        "delta": {"stop_reason": stop_reason, "stop_sequence": None},
        "usage": {"output_tokens": output_tokens}}) + "\n\n")
    yield f"{_MSG_STOP}" + '{"type":"message_stop"}' + "\n\n"


@app.post("/v1/messages")
@app.post("/v1/chat/completions")
async def messages(
    request:       Request,
    authorization: str = Header(default=None),
):
    request_id = uuid.uuid4().hex[:8]

    if not OPENCODE_API_KEY:
        return JSONResponse(status_code=401, content={"error": "Missing OPENCODE_API_KEY"})

    try:
        body = await request.json()
    except Exception:
        return JSONResponse(status_code=400, content={"error": "Invalid JSON"})

    requested_model = body.get("model", "claude-sonnet-4-6")
    mapped_model    = MODEL_MAP.get(requested_model, DEFAULT_MODEL)
    stream          = bool(body.get("stream", False))
    thinking_on     = THINKING_MODE == "enabled"

    anthropic_tools    = body.get("tools", [])
    openai_tools       = convert_tools_to_openai(anthropic_tools)
    openai_tool_choice = map_tool_choice(body.get("tool_choice"), bool(openai_tools), thinking_on)

    openai_messages = normalize_messages(body.get("messages", []))

    system_prompt = body.get("system", "")
    if system_prompt:
        if isinstance(system_prompt, list):
            system_prompt = "\n".join(
                b.get("text", "") if isinstance(b, dict) else str(b)
                for b in system_prompt
            )
        openai_messages.insert(0, {"role": "system", "content": system_prompt})

    payload = build_payload(body, openai_messages, openai_tools,
                            openai_tool_choice, mapped_model, stream)
    headers = _make_upstream_headers(OPENCODE_API_KEY)

    logger.info("[%s] %s→%s stream=%s tools=%d thinking=%s msgs=%d",
                request_id, requested_model, mapped_model,
                stream, len(openai_tools), THINKING_MODE, len(openai_messages))

    client     = request.app.state.client
    last_error = None

    for attempt in range(1, MAX_RETRIES + 1):
        if attempt > 1:
            await asyncio.sleep(0.5 * (attempt - 1))

        try:
            if stream:
                upstream_req = client.build_request(
                    "POST", UPSTREAM_URL, json=payload, headers=headers
                )
                response = await client.send(upstream_req, stream=True)

                if response.status_code == 429:
                    retry_after = float(response.headers.get("retry-after", "0"))
                    raw_err = (await response.aread()).decode(errors="ignore")
                    logger.warning("[%s] 429 — retry-after=%.0fs — body: %s",
                                   request_id, retry_after, raw_err[:400])
                    await response.aclose()
                    return JSONResponse(status_code=429, content={
                        "type": "error", "error": {
                            "type": "rate_limit_error",
                            "message": f"Upstream rate-limited. retry-after={retry_after:.0f}s. Detail: {raw_err[:200]}",
                        }}, headers={"retry-after": str(int(retry_after))} if retry_after else {})

                if response.status_code != 200:
                    raw_err = (await response.aread()).decode(errors="ignore")
                    logger.warning("[%s] upstream %d: %s",
                                   request_id, response.status_code, raw_err[:300])
                    last_error = raw_err; await response.aclose(); continue

                return StreamingResponse(
                    convert_stream_to_anthropic(response, requested_model),
                    media_type="text/event-stream",
                    headers=STREAM_RESPONSE_HEADERS,
                    background=BackgroundTask(response.aclose),
                )

            response = await client.post(UPSTREAM_URL, json=payload, headers=headers)

            if response.status_code == 429:
                retry_after = float(response.headers.get("retry-after", "0"))
                logger.warning("[%s] 429 — retry-after=%.0fs", request_id, retry_after)
                return JSONResponse(status_code=429, content={
                    "type": "error", "error": {
                        "type": "rate_limit_error",
                        "message": f"Upstream rate-limited. retry-after={retry_after:.0f}s. Detail: {response.text[:200]}",
                    }}, headers={"retry-after": str(int(retry_after))} if retry_after else {})

            if response.status_code != 200:
                logger.warning("[%s] upstream %d", request_id, response.status_code)
                last_error = response.text; continue

            raw = response.content
            if not raw or not raw.strip():
                return JSONResponse(status_code=502, content={"error": "Empty upstream response"})

            try:    parsed = _loads(raw)
            except: return JSONResponse(status_code=502, content={"error": "Invalid upstream JSON"})

            return JSONResponse(content=convert_to_anthropic(parsed, requested_model))

        except Exception as exc:
            last_error = str(exc)
            logger.exception("[%s] exception attempt %d", request_id, attempt)

    return JSONResponse(status_code=502,
                        content={"error": "Upstream failed", "details": last_error})