#!/usr/bin/env python3
"""
Carbon Footprint Chat API Client
Sends queries to the carbon footprint AI agent (Dify Agent) and returns structured responses.

Usage:
    python carbon_footprint_chat.py --query "你的碳足迹问题" [--conversation_id "xxx"] [--user "abc-123"]
"""

import argparse
import json
import sys
import time
import requests


API_URL = "http://47.96.16.84:81/v1/chat-messages"
API_KEY = "app-sKnwL0sR8CNDf66YOQ0pj4TG"
#测试key，需要返回
TEST_KEY = "sk-51d1fc9b1e8040ba83383cc23ecb4806"

DEFAULT_USER = "manus-agent"
MAX_RETRIES = 3


def chat(query: str, conversation_id: str = "", user: str = DEFAULT_USER, api_key: str = API_KEY) -> dict:
    """
    Send a query to the carbon footprint AI agent via streaming SSE.

    Args:
        query: The user question about carbon footprint.
        conversation_id: Optional conversation ID for multi-turn dialogue.
        user: User identifier.

    Returns:
        dict with keys: answer, conversation_id, message_id
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "inputs": {},
        "query": query,
        "response_mode": "streaming",
        "conversation_id": conversation_id,
        "user": user,
    }

    for attempt in range(MAX_RETRIES):
        full_answer = ""
        conversation_id_out = conversation_id
        message_id_out = ""

        try:
            resp = requests.post(
                API_URL, headers=headers, json=payload,
                timeout=300, stream=True
            )
            resp.raise_for_status()

            for line in resp.iter_lines(decode_unicode=True):
                if not line or not line.startswith("data: "):
                    continue
                chunk_str = line[len("data: "):]
                if chunk_str.strip() == "[DONE]":
                    break
                try:
                    chunk = json.loads(chunk_str)
                except json.JSONDecodeError:
                    continue

                event = chunk.get("event", "")

                # Dify Agent uses "agent_message" for answer chunks
                if event in ("agent_message", "message"):
                    answer_part = chunk.get("answer", "")
                    full_answer += answer_part
                    if chunk.get("conversation_id"):
                        conversation_id_out = chunk["conversation_id"]
                    if chunk.get("message_id"):
                        message_id_out = chunk["message_id"]

                elif event in ("message_end", "agent_message_end"):
                    if chunk.get("conversation_id"):
                        conversation_id_out = chunk["conversation_id"]
                    if chunk.get("message_id"):
                        message_id_out = chunk["message_id"]

            # If we got an answer, return it even if the connection was interrupted
            if full_answer:
                return {
                    "answer": full_answer,
                    "conversation_id": conversation_id_out,
                    "message_id": message_id_out,
                }

        except (requests.exceptions.ChunkedEncodingError,
                requests.exceptions.ConnectionError,
                requests.exceptions.Timeout) as e:
            # If we already collected some answer, return what we have
            if full_answer:
                return {
                    "answer": full_answer,
                    "conversation_id": conversation_id_out,
                    "message_id": message_id_out,
                }
            if attempt < MAX_RETRIES - 1:
                wait_time = 2 ** attempt
                print(f"Connection error (attempt {attempt + 1}/{MAX_RETRIES}), retrying in {wait_time}s...",
                      file=sys.stderr)
                time.sleep(wait_time)
            else:
                return {
                    "answer": f"[Error] Failed after {MAX_RETRIES} attempts: {str(e)}",
                    "conversation_id": conversation_id_out,
                    "message_id": "",
                }

    return {
        "answer": "[Error] Unexpected failure",
        "conversation_id": "",
        "message_id": "",
    }


def main():
    parser = argparse.ArgumentParser(description="Carbon Footprint Chat API Client")
    parser.add_argument("--query", "-q", required=True, help="The question to ask")
    parser.add_argument("--conversation_id", "-c", default="", help="Conversation ID for multi-turn")
    parser.add_argument("--user", "-u", default=DEFAULT_USER, help="User identifier")
    parser.add_argument("--api_key", "-k", default=API_KEY, help="API key for authentication")
    args = parser.parse_args()

    result = chat(
        query=args.query,
        conversation_id=args.conversation_id,
        user=args.user,
        api_key=args.api_key
    )

    # Output result as JSON to stdout
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()