import os
import json
import asyncio
import sys
import uuid
from datetime import datetime
from unittest.mock import MagicMock, AsyncMock

# Set PYTHONPATH to include the current directory
sys.path.append(os.getcwd())

from src.ircfa_interaction.schema import ChatRequest
from src.auth.model import UserBase
from src.ircfa_interaction.stream2 import stream_create
from src.ircfa_interaction.model import ChatHistory as ChatHistoryModel, Conversation as ConversationModel
from src.database.core import async_session, collector_async_session

# Token provided by USER
TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJmcmVzaCI6ZmFsc2UsImlhdCI6MTc3NTEwNzEzOS4wODk5NjksImp0aSI6InRjOU9MbndhMkxFckFPaXF1YVk4aXciLCJ0eXBlIjoiYWNjZXNzIiwic3ViIjoiMTUzZjdlMmUtZjBlYS00NDBjLTk5Y2QtN2E3ODJhZThhNjMxIiwibmJmIjoxNzc1MTA3MTM5LCJjc3JmIjoiMmRiNzI3MWEtYTU5Mi00ODU4LWJjMzAtNDY4NGU2ODI1YWNjIiwiZXhwIjoyMDkwNDY3MTM5LCJub25jZSI6ImQzYWNkMDE5YjNiNzkyNWU3ZTY0MDFiMGQyMzhmNDQxIiwidXNlcm5hbWUiOiJkZXZhZG1pbiJ9.YjmpfEFuOAvdOyp9wpPAaykdKGkHFwbNayyVPNDrw_-13-noQwY76lbns-aGF3DG35il2bbjQqNNBIAUFwcD0mDJqeWxaAQkbaL_SGLHpGU6Ul6DGNpyOE4jHxASI7B8DdiwRJkuI-zsA-4uISC8kIkyskOjv-HGFi7mwsxwScbtsWQ5d6dK9xzNCjbwYCVrskJHjEV6aXIYI1VisgbNFyx-hsvbVKKydka1Olu6iurv25ZopDgU0LwRtppALYMURcZYRm9Nrh6ZHLq5306rqUXsKwBbmzKQ8WMExfICvvk1zRPbCNG6-wJCf0E0Ft1XNuFfhbDIM-2N_zS3QCJQGw"

async def test_chat_flow(message: str, location_tag: str = None):
    print(f"\n=== Testing Gemma Tool Call Flow ===")
    print(f"User Message: {message}")
    print(f"Initial Location Tag: {location_tag}")

    # 1. Setup Mock User
    current_user = UserBase(
        user_id=str(uuid.uuid4()),
        name="Test User",
        role="admin"
    )

    # 2. Setup DB Sessions
    async with async_session() as db_session:
        async with collector_async_session() as collector_db_session:

            # 3. Prepare Conversation ID
            conv_uuid = uuid.uuid4()

            # Ensure Conversation exists in DB to prevent IntegrityError
            new_conv = ConversationModel(
                id=conv_uuid,
                conversation_id=conv_uuid,
                user_id=uuid.UUID(current_user.user_id),
                conversation_title=f"Test Gemma Tool",
                created_by=current_user.name,
                updated_by=current_user.name
            )
            db_session.add(new_conv)
            await db_session.commit()
            print(f"✅ Created dummy conversation in DB: {conv_uuid}")

            # 4. Create Request
            request_as_rag = MagicMock()
            request_as_rag.message = message
            request_as_rag.conversation_id = conv_uuid
            request_as_rag.attachments = []
            request_as_rag.equipment_id = None
            request_as_rag.model = "gemma4-e2b-llamacpp"  # Force use of Gemma model
            request_as_rag.context_state = {}

            # Mocking CommonParameters
            common = {
                "db_session": db_session,
                "page": 1,
                "items_per_page": 5
            }

            # 5. Invoke stream_create
            print("Starting stream...")
            try:
                generator = stream_create(
                    db_session=db_session,
                    collector_db_session=collector_db_session,
                    request=request_as_rag,
                    upload_type=None,
                    location_tag=location_tag,
                    current_user=current_user,
                    token=TOKEN,
                    common=common
                )

                full_response = ""
                async for chunk in generator:
                    if chunk.startswith("data: "):
                        content = chunk[6:].strip()
                        if content:
                            if content.startswith("{"):
                                try:
                                    data = json.loads(content)
                                    print(f"[Metadata]: {content}")
                                except:
                                    print(f"[Chunk]: {content}")
                            else:
                                # Normal text
                                print(content, end="", flush=True)
                                full_response += content
                    elif chunk.startswith("event: "):
                        print(f"\n[Event]: {chunk.strip()}")

                print("\n\n=== Final Response ===")
                print(full_response)

            except Exception as e:
                print(f"\n❌ Chat flow failed: {e}")
                import traceback
                traceback.print_exc()

async def main():
    # Test the specific query that was failing
    query = "Ambil data sensor PFI untuk HP FEEDWATER HEATER NO.5. Anda HARUS memanggil tool yang sesuai untuk mendapatkan data ini."
    tag = "3FW-H050"
    await test_chat_flow(query, location_tag=tag)

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