from managers.chat_manager import ChatManager
from api.telemetry import TelemetryAPI
from config.api_keys import APIKeyManager
from config.models import ModelID
from auth.register import Registration
from auth.login import open_auth_url
from api.chat.chat_api import ChatAPI, ChatConfig

import urllib3
import asyncio

urllib3.disable_warnings()

from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn

console = Console()


def select_api_key(key_manager: APIKeyManager) -> str:
    """Let user select an API key from available keys"""
    keys = key_manager.list_keys()
    if not keys:
        return None

    console.print("\n[bold cyan]Available API keys:[/]")
    for i, key in enumerate(keys):
        console.print(f"[green]{i + 1}[/]. {key}")

    while True:
        try:
            choice = int(input("\nSelect API key (number) or 0 for new login: "))
            if choice == 0:
                return None
            if 1 <= choice <= len(keys):
                return keys[choice - 1]
        except ValueError:
            return keys[0]
        print("Invalid selection. Please try again.")


async def main():
    key_manager = APIKeyManager()
    chat_manager = ChatManager()
    api_key = select_api_key(key_manager)

    # If no API key selected or available, do registration
    if api_key is None:
        _ = open_auth_url()
        firebase_token = input("Please paste your firebase token: ")
        api_key = await Registration.register_user(firebase_token.strip())
        key_manager.add_key(api_key)
        print(f"Registered successfully. API key: {api_key}")

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
    ) as progress:
        progress.add_task("Initializing Chat API...", total=None)
        tele = TelemetryAPI(api_key=api_key)
        await tele.do_telemetry()

        await chat_manager.add_chat(api_key)

        progress.stop()

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
        ) as progress:
            _ = progress.add_task("Sending message...", total=None)

            config = ChatConfig(
                model_id=ModelID.MODEL_CLAUDE_3_OPUS_20240229,
                temperature=0.7,
                max_tokens=8192,
            )

            message = [
                {"role": "user", "content": "1234567890"},
                {"role": "assistant", "content": "1234567890"},
                {"role": "user", "content": "1234567890"},
                {"role": "assistant", "content": "1234567890"},
                {"role": "user", "content": "1234567890"},
                {"role": "assistant", "content": "1234567890"},
                {"role": "user", "content": "1234567890"},
                {"role": "assistant", "content": "1234567890"},
                {"role": "user", "content": "1234567890"},
                {"role": "assistant", "content": "1234567890"},
                {"role": "user", "content": "tell me a joke"},
            ]
            first_chunk = True
            chat = await chat_manager.get_chat()

            tokens = 0

            async for chunk, token_count in chat.send_message(
                messages=message,
                config=config,
                system_prompt="You are a helpful assistant.",
                stream=True,
            ):
                if first_chunk:
                    progress.stop()
                    first_chunk = False
                    console.print("=" * 20)
                if chunk:
                    console.print(
                        chunk, end="", markup=True, highlight=True, emoji=True
                    )
                    tokens += token_count

            console.print(f"\n[bold green]Tokens: {tokens}[/]")


if __name__ == "__main__":
    asyncio.run(main())