import gradio as gr
import base64
import os
from typing import List, Optional, Dict, Any, Tuple
import anthropic
import json
from pathlib import Path


def count_tokens(
    system: str, content: str, tools: str, files: List[gr.File] | None, api_key: str
) -> Tuple[Dict, str]:
    try:
        if not content.strip():
            return None, "Error: User Message is required."

        client = anthropic.Anthropic(
            api_key=api_key.strip()
            if api_key.strip()
            else os.getenv("ANTHROPIC_API_KEY")
        )

        has_pdf = False
        has_files = bool(files and any(files))

        if has_files:
            message_content = []
            for file in files:
                if not file:
                    continue

                with open(file.name, "rb") as f:
                    file_data = base64.b64encode(f.read()).decode("utf-8")

                if file.name.lower().endswith(
                    (".png", ".jpg", ".jpeg", ".gif", ".webp")
                ):
                    message_content.append(
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": f"image/{file.name.split('.')[-1].lower()}",
                                "data": file_data,
                            },
                        }
                    )
                elif file.name.lower().endswith(".pdf"):
                    message_content.append(
                        {
                            "type": "document",
                            "source": {
                                "type": "base64",
                                "media_type": "application/pdf",
                                "data": file_data,
                            },
                        }
                    )
                    has_pdf = True

            message_content.append({"type": "text", "text": content})
            final_content = message_content
        else:
            final_content = content

        betas = ["token-counting-2024-11-01"]
        if has_pdf:
            betas.append("pdfs-2024-09-25")

        params = {
            "betas": betas,
            "model": "claude-3-5-sonnet-20241022",
            "messages": [{"role": "user", "content": final_content}],
        }

        if system and system.strip():
            params["system"] = system.strip()

        if tools and tools.strip():
            try:
                params["tools"] = json.loads(tools)
            except json.JSONDecodeError:
                return None, "Error: Invalid JSON format in Tools input"

        try:
            response = client.beta.messages.count_tokens(**params)
            return params, str(response.input_tokens)

        except anthropic.APIError as api_error:
            if "system:" in str(api_error):
                params.pop("system", None)
                response = client.beta.messages.count_tokens(**params)
                return (
                    params,
                    f"{response.input_tokens} (Note: System prompt was ignored)",
                )
            else:
                raise

    except Exception as e:
        error_msg = str(e)
        if isinstance(e, anthropic.APIError):
            error_msg = f"API Error: {e.message}"
        return None, f"Error: {error_msg}"


def create_demo(api_key: Optional[str] = None) -> gr.Blocks:
    with gr.Blocks(title="Anthropic API Token Counter") as demo:
        gr.Markdown("# Anthropic API Token Counter")
        gr.Markdown("""
        Count tokens for Anthropic API messages, including system prompts, user messages, tools, images, and PDFs. For more details, see the [official documentation](https://docs.anthropic.com/en/docs/build-with-claude/token-counting). Get your API key from the [Anthropic Console](https://console.anthropic.com/settings/keys).

        Anthropic API のメッセージのトークン数をカウントします。システムプロンプト、ユーザーメッセージ、ツール、画像、PDFを含むメッセージに対応しています。詳細は[公式ドキュメント](https://docs.anthropic.com/en/docs/build-with-claude/token-counting)を参照してください。API Keyは[Anthropic Console](https://console.anthropic.com/settings/keys)から取得できます。
        """)

        with gr.Row():
            with gr.Column():
                api_key_input = gr.Textbox(
                    label="Anthropic API Key (Required)",
                    placeholder="Enter your Anthropic API Key",
                    value="",
                    type="password",
                )

                system_input = gr.Textbox(
                    label="System Prompt (Optional)",
                    placeholder="Enter system prompt if needed",
                    lines=3,
                )

                content_input = gr.Textbox(
                    label="User Message (Required)",
                    placeholder="Enter the message you want to send to Claude",
                    lines=5,
                )

                tools_input = gr.Textbox(
                    label="Tools (Optional)",
                    placeholder="""[{
    "name": "get_weather",
    "description": "Get the current weather in a given location",
    "input_schema": {
        "type": "object",
        "properties": {
            "location": {
                "type": "string",
                "description": "The city and state, e.g. San Francisco, CA"
            }
        },
        "required": ["location"]
    }
}]""",
                    lines=8,
                )

                file_input = gr.File(
                    label="Upload Files (Optional) - Images & PDFs",
                    file_count="multiple",
                    file_types=[".jpg", ".jpeg", ".png", ".gif", ".webp", ".pdf"],
                )

                count_button = gr.Button("Count Tokens")

            with gr.Column():
                request_output = gr.JSON(
                    label="Request",
                )
                token_output = gr.Textbox(label="Token Count", interactive=False)

        count_button.click(
            fn=count_tokens,
            inputs=[
                system_input,
                content_input,
                tools_input,
                file_input,
                api_key_input,
            ],
            outputs=[request_output, token_output],
        )

        return demo


if __name__ == "__main__":
    demo = create_demo()
    demo.launch()