"""Native YunQi HTTP requests. Python 3.8+, pip install requests."""
import argparse
import base64
import json
import mimetypes
import os
from pathlib import Path

import requests


def iter_sse(response):
    data = []
    for line in response.iter_lines(decode_unicode=True):
        if line == "":
            if data:
                value = "\n".join(data)
                data = []
                yield value
            continue
        if line.startswith(":"):
            continue
        field, _, value = line.partition(":")
        if field == "data":
            data.append(value[1:] if value.startswith(" ") else value)
    if data:
        raise RuntimeError("Incomplete SSE frame at connection close")


def build_body(protocol, model, prompt, attachment=None, stream=False, mime_override=None):
    mime, encoded = None, None
    if attachment:
        path = Path(attachment)
        native_mimes = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp", ".gif": "image/gif", ".pdf": "application/pdf", ".mp4": "video/mp4", ".mpeg": "video/mpeg", ".mov": "video/mov", ".mp3": "audio/mp3", ".wav": "audio/wav", ".avi": "video/avi", ".wmv": "video/wmv", ".mpg": "video/mpg", ".csv": "text/csv", ".txt": "text/plain", ".json": "application/json", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"}
        mime = mime_override or native_mimes.get(path.suffix.lower()) or mimetypes.guess_type(str(path))[0]
        if not mime:
            raise ValueError("Cannot determine file MIME")
        encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    if protocol == "responses":
        content = [{"type": "input_text", "text": prompt}]
        if encoded:
            if mime.startswith("image/"):
                content.append({"type": "input_image", "image_url": "data:" + mime + ";base64," + encoded})
            elif mime in ["application/pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "text/plain", "text/csv", "application/json"]:
                content.append({"type": "input_file", "filename": Path(attachment).name, "file_data": "data:" + mime + ";base64," + encoded})
            else:
                raise ValueError("Responses example accepts images, PDF, DOCX, XLSX, TXT, CSV or JSON")
        return "/v1/responses", {"model": model, "input": [{"role": "user", "content": content}], "max_output_tokens": 4096, "stream": stream}
    if protocol == "messages":
        content = [{"type": "text", "text": prompt}]
        if encoded:
            if not (mime.startswith("image/") or mime == "application/pdf"):
                raise ValueError("Messages example accepts images or PDF")
            content.append({"type": "document" if mime == "application/pdf" else "image", "source": {"type": "base64", "media_type": mime, "data": encoded}})
        return "/v1/messages", {"model": model, "messages": [{"role": "user", "content": content}], "max_tokens": 8192, "stream": stream}
    if protocol == "gemini":
        from urllib.parse import quote
        parts = [{"text": prompt}]
        if encoded:
            parts.append({"inlineData": {"mimeType": mime, "data": encoded}})
        action = ":streamGenerateContent?alt=sse" if stream else ":generateContent"
        return "/v1beta/models/" + quote(model, safe="") + action, {"contents": [{"role": "user", "parts": parts}], "generationConfig": {"maxOutputTokens": 8192}}
    content = prompt
    if encoded:
        if not mime.startswith("image/"):
            raise ValueError("Chat attachment must be an image")
        content = [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": "data:" + mime + ";base64," + encoded}}]
    return "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": content}], "max_tokens": 4096, "stream": stream}


def inspect_result(protocol, data):
    if data.get("error"):
        raise RuntimeError("API error: " + json.dumps(data["error"], ensure_ascii=False))
    if protocol == "responses":
        if data.get("status") != "completed":
            raise RuntimeError("Response not completed: " + str(data.get("incomplete_details") or data.get("status")))
        text = "".join(p.get("text", "") for item in data.get("output", []) if item.get("type") == "message" for p in item.get("content", []) if p.get("type") == "output_text")
    elif protocol == "messages":
        if data.get("stop_reason") != "end_turn":
            raise RuntimeError("Message not a completed text turn: " + str(data.get("stop_reason")))
        text = "".join(x.get("text", "") for x in data.get("content", []) if x.get("type") == "text")
    elif protocol == "gemini":
        candidates = data.get("candidates", [])
        if not candidates or candidates[0].get("finishReason") != "STOP":
            raise RuntimeError("Candidate not completed: " + str([c.get("finishReason") for c in candidates]))
        text = "".join(p.get("text", "") for c in candidates for p in c.get("content", {}).get("parts", []) if not p.get("thought"))
    else:
        choices = data.get("choices", [])
        if not choices or any(c.get("finish_reason") != "stop" for c in choices):
            raise RuntimeError("Chat is truncated or requires a tool round")
        text = "".join(c.get("message", {}).get("content") or "" for c in choices)
    if not text.strip():
        raise RuntimeError("No final text returned")
    return text


def stream_text(protocol, response):
    terminal = False
    text_turn_ended = False
    text = []
    for raw in iter_sse(response):
        if raw == "[DONE]":
            if protocol == "chat":
                terminal = text_turn_ended
            continue
        event = json.loads(raw)
        kind = event.get("type")
        if kind == "error" or event.get("error") or kind in ["response.failed", "response.incomplete"]:
            raise RuntimeError("Stream error: " + json.dumps(event, ensure_ascii=False))
        delta = ""
        if protocol == "responses":
            if kind == "response.output_text.delta":
                delta = event.get("delta", "")
            if kind == "response.completed":
                terminal = event.get("response", {}).get("status") == "completed"
        elif protocol == "messages":
            if kind == "content_block_delta" and event.get("delta", {}).get("type") == "text_delta":
                delta = event["delta"].get("text", "")
            if kind == "message_delta":
                reason = event.get("delta", {}).get("stop_reason")
                if reason not in [None, "end_turn"]:
                    raise RuntimeError("Message stream ended with stop_reason=" + str(reason))
                if reason == "end_turn":
                    text_turn_ended = True
            if kind == "message_stop":
                terminal = text_turn_ended
        elif protocol == "gemini":
            for candidate in event.get("candidates", []):
                delta += "".join(p.get("text", "") for p in candidate.get("content", {}).get("parts", []) if not p.get("thought"))
                if candidate.get("finishReason"):
                    if candidate["finishReason"] != "STOP":
                        raise RuntimeError("Gemini stream ended: " + candidate["finishReason"])
                    terminal = True
        else:
            for choice in event.get("choices", []):
                delta += choice.get("delta", {}).get("content") or ""
                if choice.get("finish_reason") not in [None, "stop"]:
                    raise RuntimeError("Chat stream truncated or requires tools")
                if choice.get("finish_reason") == "stop":
                    text_turn_ended = True
        text.append(delta)
        print(delta, end="", flush=True)
    print()
    if not terminal or not "".join(text).strip():
        raise RuntimeError("Stream ended without completed text")
    return "".join(text)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("protocol", choices=["responses", "messages", "gemini", "chat"])
    prompts = parser.add_mutually_exclusive_group()
    prompts.add_argument("--prompt")
    prompts.add_argument("--prompt-file", help="UTF-8 text input; avoids shell command-length limits")
    parser.add_argument("--file")
    parser.add_argument("--mime-type", help="Explicit MIME for the real file encoding")
    parser.add_argument("--stream", action="store_true")
    parser.add_argument("--max-output-tokens", type=int, help="Positive model-supported output budget, including reasoning")
    args = parser.parse_args()
    key, model = os.environ["YUNQIAI_API_KEY"], os.environ["YUNQIAI_MODEL"]
    prompt = Path(args.prompt_file).read_text(encoding="utf-8") if args.prompt_file else args.prompt or "仅回复 READY"
    if args.mime_type and not args.file:
        parser.error("--mime-type requires --file")
    route, body = build_body(args.protocol, model, prompt, args.file, args.stream, args.mime_type)
    if args.max_output_tokens is not None:
        if args.max_output_tokens < 1:
            parser.error("--max-output-tokens must be positive")
        if args.protocol == "gemini":
            body['generationConfig']['maxOutputTokens'] = args.max_output_tokens
        else:
            body['max_output_tokens' if args.protocol == 'responses' else 'max_tokens'] = args.max_output_tokens
    headers = {"Authorization": "Bearer " + key}
    if args.protocol == "messages":
        headers = {"x-api-key": key, "anthropic-version": "2023-06-01"}
    elif args.protocol == "gemini":
        headers = {"x-goog-api-key": key}
    # No automatic retries: a disconnected POST may already be accepted.
    with requests.post("https://www.yunqiai.chat" + route, headers=headers, json=body, stream=args.stream, allow_redirects=False, timeout=(15, 300)) as response:
        response.encoding = "utf-8"
        if not 200 <= response.status_code < 300:
            raise RuntimeError("HTTP " + str(response.status_code) + ": " + response.text[:4000])
        if args.stream:
            stream_text(args.protocol, response)
        else:
            print(inspect_result(args.protocol, response.json()))


if __name__ == "__main__":
    main()