import asyncio
import json
import os
from datetime import datetime
from pathlib import Path

import matrx_ai

matrx_ai.initialize()

import rich
from matrx_utils import cleanup_async_resources, clear_terminal, to_matrx_json, vcprint

from matrx_ai.orchestrator import CompletedRequest

LOCAL_USER_ID = os.getenv("LOCAL_USER_ID")


async def register_all_tools():
    from matrx_ai.tools.registry import ToolRegistryV2
    tool_registry = ToolRegistryV2.get_instance()
    await tool_registry.load_from_database()
    vcprint(tool_registry.count, title="[EXECUTION TEST] register_tools Tools Loaded", color="blue", inline=True)



async def test_autonomous_execution(config: dict, conversation_id: str, new_conversation: bool = False, debug: bool = False):
    """Test autonomous execution that handles all tool calls automatically"""
    from matrx_ai.config.unified_config import UnifiedConfig
    from matrx_ai.orchestrator.executor import execute_ai_request
    from tests.ai.test_context import create_test_execution_context
    
    create_test_execution_context(conversation_id=conversation_id, debug=debug, new_conversation=new_conversation)

    if config.get("tools") and isinstance(config.get("tools"), list) and len(config.get("tools")) > 0:
        await register_all_tools()

    unified_config = UnifiedConfig.from_dict(config)

    try:
        completed = await execute_ai_request(
            unified_config,
            max_iterations=10,
            max_retries_per_iteration=0,
        )
        return completed
    except RuntimeError as e:
        vcprint(f"\n✗ Autonomous execution failed: {str(e)}", color="red")
        raise


def clean_up_response(response):
    last_response = response.request.config.messages.get_last_by_role("assistant")
    last_output = response.request.config.get_last_output()
    clean_response = {
        "output": last_output,
        "assistant_response": last_response,
        "config": response.request.config,
        "usage": response.total_usage,
        "metadata": response.metadata,
    }
    rich.print(clean_response)
    return clean_response


def save_results(final_result: CompletedRequest):
    output_file = Path(__file__).parent / "test_data" / "final_response.json"
    serializable_result = to_matrx_json(final_result)
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(serializable_result, f, indent=4)
    vcprint(f"\nFinal result saved to: {output_file}", color="blue")

    clean_response = clean_up_response(final_result)
    output_file = Path(__file__).parent / "test_data" / "clean_response.json"
    serializable_result = to_matrx_json(clean_response)
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(serializable_result, f, indent=4)
    vcprint(f"\nClean response saved to: {output_file}", color="blue")

    storage_data = final_result.to_storage_dict()
    storage_serializable = to_matrx_json(storage_data)
    output_file = Path(__file__).parent / "test_data" / "cx_storage_response.json"
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(storage_serializable, f, indent=4)
    vcprint(f"\nCX storage format saved to: {output_file}", color="blue")

    if os.environ.get("MATRX_CAPTURE_LLM_RESPONSES", "").lower() in ("1", "true", "yes"):
        from tests.ai.translation_tests.response_capture import CAPTURE_DIR, load_captured_responses

        captures = load_captured_responses()
        if captures:
            vcprint(f"\n{len(captures)} provider response(s) captured → {CAPTURE_DIR}", color="green")
            for cap in captures:
                meta = cap.get("capture_metadata", {})
                vcprint(
                    f"  {meta.get('provider')} / {meta.get('model')}  stream={meta.get('context', {}).get('stream')}  → {Path(cap['_filepath']).name}",
                    color="cyan",
                )



# ============================================================================
# EXAMPLE USAGE
# ============================================================================


if __name__ == "__main__":
    os.environ["MATRX_CAPTURE_LLM_RESPONSES"] = "1"
    clear_terminal()
    from tests.ai.test_context import get_test_conversation_id

    # test_user_id = get_test_user_id()
    # test_conversation_id = str(uuid.uuid4())

    openai = "gpt-5-mini"
    google = "gemini-3-flash-preview"
    model_3 = "548126f2-714a-4562-9001-0c31cbeea375"
    google_image_model = "gemini-3-pro-image-preview"
    google_image_small = "gemini-3.1-flash-image-preview"
    model_5 = "gpt-5"
    claude = "claude-sonnet-4-5-20250929"
    claude_low = "5b467c4b-80f3-420f-a516-05218907521b"
    claude_high = "claude-sonnet-4-6"
    cerebras = "gpt-oss-120b"
    xai = "grok-4-1-fast-reasoning"
    together = "openai/gpt-oss-120b"
    groq = "llama-3.3-70b-versatile"
    openai_non_reasoning = "gpt-4.1-mini-2025-04-14"
    
    
    claude_adaptive = "claude-sonnet-4-6"
    openai_reasoning = "gpt-5.2"
    google_thinking = "gemini-3.1-pro-preview"

    google_small = "gemini-3.1-flash-lite-preview"

    multiple_tools_settings = {
        "debug": True,
        "config": {
            "model": google,
            "stream": True,
            "reasoning_effort": "low",
            "reasoning_summary": "detailed",
            "tools": [
                "travel_get_location",
                "travel_get_weather",
                "travel_get_restaurants",
                "travel_get_activities",
                "travel_get_events",
                "travel_create_summary",
            ],
            "tool_choice": "auto",
            "parallel_tool_calls": True,
            "internal_web_search": False,
            "internal_url_context": False,
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant. Think through your response step by step and use the many tools you have available to you. You have tools to get you the user's location, the weather, restaurants, activities, and events in the area.",
                },
                {
                    "role": "user",
                    "content": "Hi",
                },
                {
                    "role": "assistant",
                    "content": "Hello! How may I assist you today?",
                },
                {
                    "role": "user",
                    "content": "Hello! Can you give me some recommendations for today? restaurants and events in the area?",
                },
            ],
        },
    }

    image_generation_settings = {
        "ai_model_id": google_image_model,
        "debug": True,
        "config": {
            # "max_output_tokens": 4444,
            "stream": True,
            "reasoning_effort": "high",
            "reasoning_summary": "always",
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant.",
                },
                {
                    "role": "user",
                    "content": "Can you generate an image of a cat?",
                },
            ],
        },
    }

    complex_settings = {
        "ai_model_id": cerebras,
        "debug": False,
        "config": {
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant. Think through your response step by step before proceeding with a response. Always search the web to ensure you include recent and relevant facts in your response. You can check the user's location with a tool and provide whatever they need.",
                },
                {
                    "role": "user",
                    "content": "Hello! Can you help me? What is the latest us news?",
                },
            ],
            "temperature": 1,
            "max_output_tokens": 4096,
            "top_p": 1,
            "top_k": 50,
            "frequency_penalty": 0,
            "presence_penalty": 0,
            "stream": True,
            "response_format": "text",
            "store": True,
            "stop_sequences": [],
            "tools": [
                "travel_get_location",
                "travel_get_weather",
                "travel_get_restaurants",
                "travel_get_activities",
                "travel_get_events",
                "travel_create_summary",
            ],
            "tool_choice": "auto",
            "parallel_tool_calls": False,
            "image_urls": False,
            "file_urls": False,
            "internal_web_search": True,
            "youtube_videos": False,
            "size": "1024x1024",
            "quality": "standard",
            "count": 1,
            "audio_voice": "alloy",
            "audio_format": "mp3",
            "reasoning_effort": "low",
            "verbosity": "medium",
            "reasoning_summary": "auto",
            "modalities": {},
            "preset_name": "default",
            "system_message_override": None,
            "reasoning": {},
        },
    }

    simple_chat_settings = {
        "debug": True,
        "config": {
            "model": "gpt-5-mini",
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant.",
                },
                {
                    "role": "user",
                    "content": "Hello! How are you?",
                },
            ],
            "stream": True,
        },
    }

    single_tool_settings = {
        "debug": False,
        "config": {
            # "max_output_tokens": 4444,
            "model": google,
            "stream": True,
            "reasoning_effort": "low",
            "reasoning_summary": "detailed",
            "tools": [
                "news_get_headlines",
            ],
            "tool_choice": "auto",
            "parallel_tool_calls": True,
            "internal_web_search": False,
            "internal_url_context": False,
            "messages": [
                {
                    "role": "system",
                    "content": f"You're a helpful assistant. Think through your response step by step and use the many tools you have available to you. You have tools to get the latest news headlines. Today's Date: {datetime.now().strftime('%Y-%m-%d')}",
                },
                {
                    "role": "user",
                    "content": "Hi",
                },
                {
                    "role": "assistant",
                    "content": "Hello! How may I assist you today?",
                },
                {
                    "role": "user",
                    "content": "Hello! Can you give me some of the latest news headlines? Just for the US is good and I want them up to  date for today.",
                },
            ],
        },
    }

    # Same as single_tool_settings but using structured system_instruction.
    # Date is auto-injected by SystemInstruction (no f-string hack needed).
    # No {"role": "system", ...} message in the messages array.
    single_tool_settings_v2 = {
        "debug": False,
        "config": {
            "model": google,
            "stream": True,
            "reasoning_effort": "low",
            "reasoning_summary": "detailed",
            "system_instruction": {
                "content": "You're a helpful assistant. Think through your response step by step and use the many tools you have available to you. You have tools to get the latest news headlines.",
                "include_date": True,
            },
            "tools": [
                "news_get_headlines",
            ],
            "tool_choice": "auto",
            "parallel_tool_calls": True,
            "internal_web_search": False,
            "internal_url_context": False,
            "messages": [
                {
                    "role": "user",
                    "content": "Hi",
                },
                {
                    "role": "assistant",
                    "content": "Hello! How may I assist you today?",
                },
                {
                    "role": "user",
                    "content": "Hello! Can you give me some of the latest news headlines? Just for the US is good and I want them up to date for today.",
                },
            ],
        },
    }

    image_input_settings = {
        "debug": True,
        "config": {
            "model": openai_non_reasoning,
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant.",
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_text",
                            "text": "what is in this image? be specific and detailed.",
                        },
                        {
                            "type": "input_image",
                            "image_url": f"https://txzxabzwovsujtloxrus.supabase.co/storage/v1/object/public/user-public-assets/user-{LOCAL_USER_ID}/mcp-logo.png",
                        },
                    ],
                },
            ],
            "stream": True,
        },
    }

    document_input_settings = {
        "debug": True,
        "config": {
            "model": openai_non_reasoning,
            "messages": [
                {
                    "role": "system",
                    "content": "You're a document analyzer. Only responsd if you can clearly see the document and it's content.",
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_text",
                            "text": "what is in this document? be specific and detailed.",
                        },
                        {
                            "type": "input_document",
                            "url": f"https://txzxabzwovsujtloxrus.supabase.co/storage/v1/object/public/user-public-assets/user-{LOCAL_USER_ID}/Ads%20Conv%20Analysis.pdf",
                        },
                    ],
                },
            ],
            "stream": True,
        },
    }

    youtube_url_settings = {
        "debug": True,
        "config": {
            "model": google,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "youtube_video",
                            "youtube_url": "https://www.youtube.com/watch?v=9hE5-98ZeCg",
                        },
                        {
                            "type": "text",
                            "text": "Please summarize this video in 3 sentences.",
                        },
                    ],
                },
            ],
            "stream": True,
        },
    }

    # Audio URL for testing
    audio_url = "https://txzxabzwovsujtloxrus.supabase.co/storage/v1/object/public/Audio/Audio%20Files/My%20Recordings/small_sample_audio.m4a"

    # Test 1: Use Groq transcription to convert audio to text, then send to Google
    audio_transcription_test_settings = {
        "debug": True,
        "config": {
            "model": openai_non_reasoning,
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant that analyzes audio content.",
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "audio",
                            "url": audio_url,
                            # "auto_transcribe": True,  # Transcribe with Groq first
                            # "transcription_model": "whisper-large-v3-turbo", # Optional
                            # "transcription_language": "en", # Optional
                        },
                        {
                            "type": "text",
                            "text": "Please summarize what was said in the audio in 2-3 sentences.",
                        },
                    ],
                },
            ],
            "stream": True,
        },
    }

    # Test 2: Send audio directly to Google (Google can handle audio natively)
    audio_direct_google_settings = {
        "debug": True,
        "config": {
            "model": google,
            "messages": [
                {
                    "role": "system",
                    "content": "You're a helpful assistant that analyzes audio content.",
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "audio",
                            "url": audio_url,
                            "auto_transcribe": False,  # Send directly to Google
                        },
                        {
                            "type": "text",
                            "text": "Please summarize what was said in the audio in 2-3 sentences.",
                        },
                    ],
                },
            ],
            "stream": True,
        },
    }

    # options: simple_chat_settings, single_tool_settings_v2, image_generation_settings, complex_settings,
    # document_input_settings, youtube_url_settings, audio_transcription_test_settings, audio_direct_google_settings
    settings_to_use = single_tool_settings
    
    # claude_adaptive
    # google_thinking    | google_small
    # openai_reasoning
    
    
    # cerebras xai together groq
    
    settings_to_use["config"]["model"] = "gpt-5.4"
    new_conversation = True
    # settings_to_use["debug"] = True


    async def run():
        conversation_id = get_test_conversation_id()
        final_result: CompletedRequest = await test_autonomous_execution(settings_to_use["config"], conversation_id, new_conversation=new_conversation, debug=False)
        save_results(final_result)
        
    asyncio.run(run())
    cleanup_async_resources()
