import asyncio
import json

import httpx
from fastapi import APIRouter, Depends, Request
from fastapi.responses import Response, StreamingResponse

from letta.log import get_logger
from letta.server.rest_api.dependencies import HeaderParams, get_headers, get_letta_server
from letta.server.rest_api.proxy_helpers import (
    build_response_from_chunks,
    check_for_duplicate_message,
    extract_assistant_message,
    extract_user_messages,
    get_or_create_claude_code_agent,
    inject_memory_context,
    is_topic_detection_response,
    persist_messages_background,
    prepare_headers,
)
from letta.server.server import SyncServer

logger = get_logger(__name__)

router = APIRouter(prefix="/zai", tags=["zai"])

ZAI_API_BASE = "https://api.z.ai/api/anthropic"
PROXY_NAME = "Z.ai Proxy"


@router.api_route("/v1/messages", methods=["POST"], operation_id="zai_messages_proxy", include_in_schema=False)
async def zai_messages_proxy(
    request: Request,
    server: SyncServer = Depends(get_letta_server),
    headers: HeaderParams = Depends(get_headers),
):
    """
    Proxy endpoint for Z.ai Messages API.

    This endpoint forwards requests to the Z.ai API, allowing Claude Code CLI
    to use Letta as a proxy by configuring anthropic_base_url.

    Usage in Claude Code CLI settings.json:
    {
        "env": {
            "ANTHROPIC_BASE_URL": "http://localhost:3000/v1/zai"
        }
    }
    """
    # Get the request body
    body = await request.body()

    logger.info(f"[{PROXY_NAME}] Proxying request to Z.ai Messages API: {ZAI_API_BASE}/v1/messages")
    logger.debug(f"[{PROXY_NAME}] Request body preview: {body[:200]}...")

    actor = await server.user_manager.get_actor_or_default_async(headers.actor_id)

    # Extract all user messages from request
    all_user_messages = extract_user_messages(body)

    # Only capture the LAST user message (the new one the user just sent)
    # Claude Code sends full conversation history, but we only want to persist the new message
    user_messages = [all_user_messages[-1]] if all_user_messages else []

    # Filter out system/metadata requests and policy specs
    user_messages = [s for s in user_messages if not s.startswith("<system-reminder>") and not s.startswith("<policy_spec>")]
    if not user_messages:
        logger.debug(f"[{PROXY_NAME}] Skipping capture/memory for this turn")

    zai_headers = prepare_headers(request, PROXY_NAME, use_bearer_auth=True)
    if not zai_headers:
        logger.error(f"[{PROXY_NAME}] No Anthropic API key found in headers or settings")
        return Response(
            content='{"error": {"type": "authentication_error", "message": "Anthropic API key required. Pass via anthropic-api-key or x-api-key header."}}',
            status_code=401,
            media_type="application/json",
        )

    # Check if this is a streaming request
    try:
        import json

        request_data = json.loads(body)
        is_streaming = request_data.get("stream", False)
        model_name = request_data.get("model")
        # Extract and remove project_id (internal use only, not for Z.ai API)
        project_id = request_data.pop("project_id", None)
        logger.debug(f"[{PROXY_NAME}] Request is streaming: {is_streaming}")
        logger.debug(f"[{PROXY_NAME}] Model: {model_name}")
        logger.debug(f"[{PROXY_NAME}] Project ID: {project_id}")
    except Exception as e:
        logger.warning(f"[{PROXY_NAME}] Failed to parse request body: {e}")
        is_streaming = False
        model_name = None
        project_id = None

    # Get or create agent for Claude Code session (skip for system requests)
    # Note: Agent lookup and memory search are blocking operations before forwarding.
    # Message persistence happens in the background after the response is returned.
    agent = None
    try:
        agent = await get_or_create_claude_code_agent(
            server=server,
            actor=actor,
            project_id=project_id,
        )
        logger.debug(f"[{PROXY_NAME}] Using agent ID: {agent.id}")
    except Exception as e:
        logger.error(f"[{PROXY_NAME}] Failed to get/create agent: {e}")

    # Inject memory context into request (skip for system requests)
    # TODO: Optimize - skip memory injection on subsequent messages in same session
    # TODO: Add caching layer to avoid duplicate memory searches
    modified_body = body
    if agent and request_data:
        modified_request_data = await inject_memory_context(
            server=server,
            agent=agent,
            actor=actor,
            request_data=request_data,
            proxy_name=PROXY_NAME,
        )
        # Re-encode the modified request
        import json

        modified_body = json.dumps(modified_request_data).encode("utf-8")

    # Forward the request to Z.ai API (preserve query params like ?beta=true)
    # Note: For streaming, we create the client outside the generator to keep it alive
    zai_url = f"{ZAI_API_BASE}/v1/messages"
    if request.url.query:
        zai_url = f"{zai_url}?{request.url.query}"

    if is_streaming:
        # Handle streaming response
        collected_chunks = []

        async def stream_response():
            # Create client inside the generator so it stays alive during streaming
            async with httpx.AsyncClient(timeout=300.0) as client:
                async with client.stream(
                    "POST",
                    zai_url,
                    headers=zai_headers,
                    content=modified_body,
                ) as response:
                    async for chunk in response.aiter_bytes():
                        collected_chunks.append(chunk)
                        yield chunk

                # After streaming is complete, extract and log assistant message
                assistant_message = build_response_from_chunks(collected_chunks)
                if user_messages and assistant_message:
                    logger.info("=" * 70)
                    logger.info("📨 CAPTURED USER MESSAGE:")
                    for i, user_message in enumerate(user_messages):
                        logger.info(f"  {i}: {user_message[:200]}{'...' if len(user_message) > 200 else ''}")
                    logger.info("=" * 70)
                    logger.info("🤖 CAPTURED ASSISTANT RESPONSE (streaming):")
                    logger.info(f"  {assistant_message[:200]}{'...' if len(assistant_message) > 200 else ''}")
                    logger.info("=" * 70)

                    # Skip persisting topic detection responses (metadata, not conversation)
                    if is_topic_detection_response(assistant_message):
                        logger.debug(f"[{PROXY_NAME}] Skipping persistence - topic detection response")
                    # Persist messages to database (non-blocking, skip for system requests)
                    elif agent:
                        # Check for duplicate user messages before creating background task
                        # This prevents race conditions where multiple requests persist the same message
                        user_messages_to_persist = await check_for_duplicate_message(server, agent, actor, user_messages, PROXY_NAME)

                        asyncio.create_task(
                            persist_messages_background(
                                server=server,
                                agent=agent,
                                actor=actor,
                                user_messages=user_messages_to_persist,
                                assistant_message=assistant_message,
                                model_name=model_name,
                                proxy_name=PROXY_NAME,
                            )
                        )

        return StreamingResponse(
            stream_response(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            },
        )

    # Non-streaming path
    async with httpx.AsyncClient(timeout=300.0) as client:
        try:
            # Handle non-streaming response
            response = await client.post(
                zai_url,
                headers=zai_headers,
                content=modified_body,
            )

            logger.info(f"Successfully proxied request, status: {response.status_code}")

            # Extract and log assistant message
            if response.status_code == 200:
                try:
                    import json

                    response_data = json.loads(response.content)
                    assistant_message = extract_assistant_message(response_data)
                    if assistant_message:
                        logger.info("=" * 70)
                        logger.info("🤖 CAPTURED ASSISTANT RESPONSE:")
                        logger.info(f"  {assistant_message[:500]}{'...' if len(assistant_message) > 500 else ''}")
                        logger.info("=" * 70)

                        # Skip persisting topic detection responses (metadata, not conversation)
                        if is_topic_detection_response(assistant_message):
                            logger.debug(f"[{PROXY_NAME}] Skipping persistence - topic detection response")
                        # Persist messages to database (non-blocking)
                        elif agent:
                            # Check for duplicate user messages before creating background task
                            user_messages_to_persist = await check_for_duplicate_message(server, agent, actor, user_messages, PROXY_NAME)

                            asyncio.create_task(
                                persist_messages_background(
                                    server=server,
                                    agent=agent,
                                    actor=actor,
                                    user_messages=user_messages_to_persist,
                                    assistant_message=assistant_message,
                                    model_name=model_name,
                                    proxy_name=PROXY_NAME,
                                )
                            )
                except Exception as e:
                    logger.warning(f"[{PROXY_NAME}] Failed to extract assistant response for logging: {e}")

            return Response(
                content=response.content,
                status_code=response.status_code,
                media_type=response.headers.get("content-type", "application/json"),
                headers={
                    k: v
                    for k, v in response.headers.items()
                    if k.lower() not in ["content-encoding", "content-length", "transfer-encoding", "connection"]
                },
            )

        except httpx.HTTPError as e:
            logger.error(f"[{PROXY_NAME}] Error proxying request to Z.ai API: {e}")
            return Response(
                content=f'{{"error": {{"type": "api_error", "message": "Failed to proxy request to Z.ai API: {str(e)}"}}}}',
                status_code=500,
                media_type="application/json",
            )


@router.api_route(
    "/v1/{endpoint:path}",
    methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
    operation_id="zai_catchall_proxy",
    include_in_schema=False,
)
async def zai_catchall_proxy(
    endpoint: str,
    request: Request,
    server: SyncServer = Depends(get_letta_server),
    headers: HeaderParams = Depends(get_headers),
):
    """
    Catch-all proxy for other Z.ai API endpoints.

    This forwards all other requests (like /v1/messages/count_tokens) directly to Z.ai
    without message capture or memory injection.
    """
    # Skip the /v1/messages endpoint (handled by specific route)
    if endpoint == "messages" and request.method == "POST":
        # This should be handled by the specific route, but just in case return error
        return Response(
            content='{"error": {"type": "routing_error", "message": "Use specific /v1/messages endpoint"}}',
            status_code=500,
            media_type="application/json",
        )

    # Get the request body
    body = await request.body()

    # Reconstruct the full path
    path = f"v1/{endpoint}"

    logger.info(f"[{PROXY_NAME}] Proxying catch-all request: {request.method} /{path}")

    zai_headers = prepare_headers(request, PROXY_NAME, use_bearer_auth=True)
    if not zai_headers:
        logger.error(f"[{PROXY_NAME}] No Anthropic API key found in headers or settings")
        return Response(
            content='{"error": {"type": "authentication_error", "message": "Anthropic API key required"}}',
            status_code=401,
            media_type="application/json",
        )

    # Forward the request to Z.ai API
    async with httpx.AsyncClient(timeout=300.0) as client:
        try:
            response = await client.request(
                method=request.method,
                url=f"{ZAI_API_BASE}/{path}",
                headers=zai_headers,
                content=body if body else None,
            )

            return Response(
                content=response.content,
                status_code=response.status_code,
                media_type=response.headers.get("content-type", "application/json"),
                headers={
                    k: v
                    for k, v in response.headers.items()
                    if k.lower() not in ["content-encoding", "content-length", "transfer-encoding", "connection"]
                },
            )

        except httpx.HTTPError as e:
            logger.error(f"[{PROXY_NAME}] Error proxying catch-all request to Z.ai API: {e}")
            return Response(
                content=f'{{"error": {{"type": "api_error", "message": "Failed to proxy request to Z.ai API: {str(e)}"}}}}',
                status_code=500,
                media_type="application/json",
            )
