from __future__ import annotations

import asyncio
import json
import uuid
from functools import partial
from typing import Any, Mapping, Optional

from .auth import resolve_bearer_token
from .causality import create_call_lineage, set_current_call_id
from .endpoints import _build_orchestra_send_endpoint
from .headers import _apply_response_ids, _header_get
from .parsing import _extract_text_from_result, _parse_sse_events
from .transport import _request_raw


def call_orchestra_agent(
    agent_id: Optional[str] = None,
    input_message: Optional[str] = None,
    *,
    agent_url: str = "https://node.a2aorchestra.com/api/v1/agents/sendmessage",
    stream: bool = False,
    context_id: Optional[str] = None,
    parts: Optional[list[dict[str, Any]]] = None,
    parent_call_id: Optional[str] = None,
    user_token: Optional[str] = None,
    headers: Optional[Mapping[str, str]] = None,
    timeout: float = 60.0,
    debug: bool = False,
) -> dict[str, Any]:
    """
    High-level agent call helper.

    Builds JSON-RPC envelope automatically, forwards tracking IDs, and handles
    both regular JSON and `text/event-stream` responses.
    """
    message_parts = list(parts or [])
    if input_message is not None and not message_parts:
        message_parts = [{"kind": "text", "text": input_message}]

    payload = {
        "id": str(uuid.uuid4()),
        "jsonrpc": "2.0",
        "method": "message/stream" if stream else "message/send",
        "params": {
            "message": {
                "kind": "message",
                "messageId": str(uuid.uuid4()),
                "contextId": context_id,
                "role": "user",
                "parts": message_parts,
            }
        },
    }

    call_id, inferred_parent_call_id, sequence_number = create_call_lineage()
    effective_parent_call_id = parent_call_id or inferred_parent_call_id

    merged_headers: dict[str, str] = dict(headers or {})
    merged_headers["Accept"] = "text/event-stream" if stream else "application/json"
    merged_headers["X-Call-ID"] = call_id
    merged_headers["X-Sequence-Number"] = str(sequence_number)
    if effective_parent_call_id:
        merged_headers["X-Parent-Call-ID"] = str(effective_parent_call_id)
    bearer_token = resolve_bearer_token(user_token)
    if bearer_token:
        merged_headers["Authorization"] = f"Bearer {bearer_token}"

    endpoint = _build_orchestra_send_endpoint(agent_url, agent_id)

    debug_headers = dict(merged_headers)
    if "Authorization" in debug_headers:
        debug_headers["Authorization"] = "Bearer ***"
    debug_data = {
        "request_endpoint": endpoint,
        "request_headers": debug_headers,
        "request_payload": payload,
    }

    raw = _request_raw(
        endpoint,
        payload=payload,
        method="POST",
        headers=merged_headers,
        timeout=timeout,
    )

    response_headers = raw["headers"]
    ids = _apply_response_ids(response_headers)
    set_current_call_id(call_id)
    body = raw["body"]
    content_type = (_header_get(response_headers, "Content-Type") or "").lower()

    # Non-2xx
    if not raw["ok"]:
        text = body.decode("utf-8", errors="replace") if body else ""
        error_data: Any = None
        if text:
            try:
                error_data = json.loads(text)
            except Exception:
                error_data = {"error": text}
        return {
            "success": False,
            "status_code": raw["status"],
            "error": error_data,
            "trace_id": ids["trace_id"],
            "session_id": ids["session_id"],
            "request_id": ids["request_id"],
            "call_id": call_id,
            "parent_call_id": effective_parent_call_id,
            "sequence_number": sequence_number,
            **(debug_data if debug else {}),
        }

    # Streaming response
    if stream:
        raw_text = body.decode("utf-8", errors="replace") if body else ""
        if "text/event-stream" not in content_type:
            return {
                "success": False,
                "status_code": raw["status"],
                "error": {"error": f"Expected text/event-stream, got {content_type}"},
                "raw_text": raw_text,
                "trace_id": ids["trace_id"],
                "session_id": ids["session_id"],
                "request_id": ids["request_id"],
                "call_id": call_id,
                "parent_call_id": effective_parent_call_id,
                "sequence_number": sequence_number,
                **(debug_data if debug else {}),
            }

        events = _parse_sse_events(raw_text)
        texts: list[str] = []
        last_result: dict[str, Any] = {}
        for event in events:
            result = event.get("result", {})
            if isinstance(result, dict):
                last_result = result
                text_piece = _extract_text_from_result(result)
                if text_piece:
                    texts.append(text_piece)

        return {
            "success": True,
            "status_code": raw["status"],
            "text": "".join(texts),
            "events": events,
            "context_id": last_result.get("contextId") if isinstance(last_result, dict) else None,
            "agent_request_id": last_result.get("requestId") if isinstance(last_result, dict) else None,
            "trace_id": ids["trace_id"],
            "session_id": ids["session_id"],
            "request_id": ids["request_id"],
            "call_id": call_id,
            "parent_call_id": effective_parent_call_id,
            "sequence_number": sequence_number,
            **(debug_data if debug else {}),
        }

    # Standard JSON response
    data: Any = None
    body_text = body.decode("utf-8", errors="replace") if body else ""
    if body_text:
        try:
            data = json.loads(body_text)
        except Exception:
            data = body_text

    result = data.get("result", {}) if isinstance(data, dict) else {}
    text = _extract_text_from_result(result) if isinstance(result, dict) else ""

    return {
        "success": True,
        "status_code": raw["status"],
        "text": text,
        "data": data,
        "context_id": result.get("contextId") if isinstance(result, dict) else None,
        "agent_request_id": result.get("requestId") if isinstance(result, dict) else None,
        "trace_id": ids["trace_id"],
        "session_id": ids["session_id"],
        "request_id": ids["request_id"],
        "call_id": call_id,
        "parent_call_id": effective_parent_call_id,
        "sequence_number": sequence_number,
        **(debug_data if debug else {}),
    }


async def call_orchestra_agent_async(
    agent_id: Optional[str] = None,
    input_message: Optional[str] = None,
    *,
    agent_url: str = "https://node.a2aorchestra.com/api/v1/agents/sendmessage",
    stream: bool = False,
    context_id: Optional[str] = None,
    parts: Optional[list[dict[str, Any]]] = None,
    parent_call_id: Optional[str] = None,
    user_token: Optional[str] = None,
    headers: Optional[Mapping[str, str]] = None,
    timeout: float = 60.0,
    debug: bool = False,
) -> dict[str, Any]:
    """Async wrapper for call_orchestra_agent."""
    fn = partial(
        call_orchestra_agent,
        agent_id,
        input_message,
        agent_url=agent_url,
        stream=stream,
        context_id=context_id,
        parts=parts,
        parent_call_id=parent_call_id,
        user_token=user_token,
        headers=headers,
        timeout=timeout,
        debug=debug,
    )
    return await asyncio.to_thread(fn)
