import os
import sys
import tempfile
import shutil
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))

from pygeai.templates.engine import TemplateEngine
from pygeai.assistant.managers import AssistantManager
from pygeai.assistant.rag.clients import RAGAssistantClient
from pygeai.assistant.rag.models import (
    Search,
    RetrieverOptions,
    SearchOptions,
    ChunkOptions,
    IndexOptions,
    RAGAssistant,
)
from pygeai.core.models import WelcomeData, LlmSettings


def test_chat_rag_streaming_template():
    """
    Test chat-rag-streaming template by:
    1. Creating a RAG assistant with a document
    2. Generating chat interface
    3. Testing programmatic chat (simulating user input)
    """
    test_dir = tempfile.mkdtemp(prefix="pygeai_chat_rag_test_")
    assistant_name = f"test_chat_rag_{os.urandom(4).hex()}"
    assistant_id = None

    print(f"Testing in: {test_dir}")

    try:
        # Step 1: Create RAG assistant with document
        print("[STEP 1] Creating RAG assistant...")
        manager = AssistantManager()

        llm_options = LlmSettings(
            cache=False,
            temperature=0.7,
            max_tokens=500,
            model_name="gpt-4",
            n=1,
            presence_penalty=0,
            frequency_penalty=0,
            provider="OpenAI",
            stream=False,
            top_p=1.0,
            type=None,
            verbose=False,
        )

        retriever_options = RetrieverOptions(type="vectorStore")

        search_options = SearchOptions(
            history_count=5,
            llm=llm_options,
            search=Search(
                k=3,
                return_source_documents=True,
                score_threshold=0.7,
                prompt="Use the following context to answer the question:\n\nContext: {context}\n\nQuestion: {question}",
                template="",
            ),
            retriever=retriever_options,
        )

        chunk_options = ChunkOptions(chunk_size=1000, chunk_overlap=200)

        index_options = IndexOptions(chunks=chunk_options)

        welcome_data = WelcomeData(
            title=f"{assistant_name} Assistant",
            description="Test RAG assistant for chat",
            features=[],
            examples_prompt=[],
        )

        rag_assistant = RAGAssistant(
            name=assistant_name,
            description="Test RAG assistant for chat",
            search_options=search_options,
            index_options=index_options,
            welcome_data=welcome_data,
        )

        response = manager.create_assistant(rag_assistant)
        assistant_id = response

        if not assistant_id:
            print("[FAIL] Could not create RAG assistant")
            sys.exit(1)

        print(f"[PASS] RAG assistant created: {assistant_id}")

        # Upload a test document
        print("[STEP 2] Uploading test document...")
        doc_path = Path(test_dir) / "test_doc.txt"
        doc_path.write_text(
            "This is a test document about Python programming. "
            "Python is a high-level programming language known for simplicity. "
            "It supports multiple programming paradigms including procedural, "
            "object-oriented, and functional programming."
        )

        rag_client = RAGAssistantClient()
        rag_client.upload_document(
            name=assistant_name, file_path=str(doc_path), content_type="text/plain"
        )

        print("[PASS] Document uploaded successfully")

        # Step 2: Generate chat template
        print("[STEP 3] Generating chat-rag-streaming template...")
        engine = TemplateEngine()

        params = {
            "name": assistant_name,
            "description": "Test streaming chat with RAG",
        }

        generated_files = engine.generate("chat-rag-streaming", Path(test_dir), params)

        assert len(generated_files) == 2, (
            f"Expected 2 files, got {len(generated_files)}"
        )
        print("[PASS] Chat template generated successfully")

        chat_script = Path(test_dir) / f"chat_{assistant_name}.py"
        readme = Path(test_dir) / "README.md"

        assert chat_script.exists(), "chat script not found"
        assert readme.exists(), "README not found"

        print("[PASS] All expected files exist")

        # Step 3: Test the chat script programmatically
        print("[STEP 4] Testing chat functionality...")

        # Instead of running interactive script, import and test the ChatClient directly
        from pygeai.chat.clients import ChatClient
        from pygeai.chat.settings import LLM_SETTINGS

        client = ChatClient()

        messages = [{"role": "user", "content": "What is Python?"}]

        response_stream = client.chat_completion(
            model=f"saia:search:{assistant_name}",
            messages=messages,
            stream=True,
            **LLM_SETTINGS,
        )

        full_response = ""
        for chunk in response_stream:
            full_response += chunk

        if not full_response:
            print("[FAIL] Chat returned empty response")
            sys.exit(1)

        print(f"[PASS] Chat response received (length: {len(full_response)})")
        print(f"Response preview: {full_response[:100]}...")

        # Verify response mentions Python
        if "python" not in full_response.lower():
            print("[WARN] Response may not be grounded in uploaded document")
        else:
            print("[PASS] Response appears grounded in document")

        print(
            "\n[PASS] ALL TESTS PASSED - chat-rag-streaming template is working correctly!"
        )

    except Exception as e:
        print(f"[FAIL] Test failed with error: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)

    finally:
        print("\nCleaning up...")
        if assistant_id:
            try:
                manager = AssistantManager()
                manager.delete_assistant(assistant_id=assistant_id)
                print(f"Deleted test assistant: {assistant_id}")
            except Exception as e:
                print(f"Warning: Could not delete assistant: {e}")

        if os.path.exists(test_dir):
            shutil.rmtree(test_dir)
            print(f"Deleted test directory: {test_dir}")


if __name__ == "__main__":
    test_chat_rag_streaming_template()
