import queue
import threading
import logging
import os
from django.db import transaction, IntegrityError
from django.http import HttpResponse
from dotenv import load_dotenv
from rest_framework.decorators import action
from rest_framework import status, viewsets
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework.views import APIView
from asgiref.sync import async_to_sync
from django.shortcuts import get_object_or_404
from channels.layers import get_channel_layer
from django.http import StreamingHttpResponse
from .embeddings_service import publish_to_neo4j
import requests
from haystack.dataclasses import (
    ChatMessage as HaystackChatMessage,
    ToolCall,
    ToolCallResult,
    TextContent,
)
from .models import ChatMessage, ChatSession, KnowledgeBase
from .rag_service import (
    get_rag_components,
    get_rag_response_dynamic,
    initialize_rag_components,
)
from .serializers import ChatSessionSerializer, KnowledgeBaseSerializer

logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()
SUNBIRD_AUTH_TOKEN = os.getenv("SUNBIRD_AUTH_TOKEN")
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")

language_map = {"eng": "eng", "lug": "lug", "teo": "teo", "nyn": "nyn", "ach": "ach"}


def translate_text(source_lang, target_lang, text, headers):
    """Translate text using Sunbird API."""
    if not SUNBIRD_AUTH_TOKEN:
        raise ValueError("Sunbird API token not configured.")

    url = "https://api.sunbird.ai/tasks/nllb_translate"
    payload = {
        "source_language": source_lang,
        "target_language": target_lang,
        "text": text,
    }
    try:
        response = requests.post(url, json=payload, headers=headers, timeout=10)
        response.raise_for_status()
        translated_text = response.json()["output"].get("translated_text")
        return translated_text or text
    except requests.exceptions.RequestException as e:
        logger.error(f"Sunbird API request failed: {e}")
        raise ValueError(f"Translation service unavailable: {e}")


class ChatSessionViewSet(viewsets.ModelViewSet):
    serializer_class = ChatSessionSerializer
    # permission_classes = [IsAuthenticated]

    def get_queryset(self):
        return self.request.user.chat_sessions.all()

    def perform_create(self, serializer):
        with transaction.atomic():
            serializer.save(user=self.request.user)
            # Enforce max 10 sessions
            user_sessions = self.request.user.chat_sessions.order_by("-updated_at")
            if user_sessions.count() > 10:
                for session in user_sessions[10:]:
                    logger.info(f"Deleting old chat session {session.id}")
                    session.delete()


class ChatView(APIView):
    # permission_classes = [IsAuthenticated]

    def post(self, request, *args, **kwargs):
        query = request.data.get("query")
        language = request.data.get("language")
        session_id = request.data.get("session_id")

        if not all([query, language, session_id]):
            return Response(
                {"error": "Query, language, and session_id are required"},
                status=status.HTTP_400_BAD_REQUEST,
            )

        if language not in language_map:
            return Response(
                {"error": f"Unsupported language code: {language}"},
                status=status.HTTP_400_BAD_REQUEST,
            )

        # Ensure RAG components
        rag_components = get_rag_components()
        if not rag_components.get("llm"):
            try:
                initialize_rag_components(MISTRAL_API_KEY)
                rag_components = get_rag_components()
            except Exception as e:
                logger.error(f"Failed to initialize RAG components: {e}")
                return Response(
                    {"error": "AI service unavailable."},
                    status=status.HTTP_503_SERVICE_UNAVAILABLE,
                )

        # Get chat session
        try:
            session = request.user.chat_sessions.get(id=session_id)
        except ChatSession.DoesNotExist:
            return Response(
                {"error": "Chat session not found."}, status=status.HTTP_404_NOT_FOUND
            )

        sunbird_headers = {
            "Authorization": f"Bearer {SUNBIRD_AUTH_TOKEN}",
            "Content-Type": "application/json",
        }

        try:
            # Save user message
            ChatMessage.objects.create(
                session=session, role="user", content=query, language=language
            )

            # Prepare chat history for LLM
            messages_for_llm = []
            for msg in session.messages.order_by("timestamp"):
                translated_content = msg.content
                if msg.language != "eng":
                    translated_content = translate_text(
                        msg.language, "eng", msg.content, sunbird_headers
                    )

                if msg.role == "user":
                    messages_for_llm.append(
                        HaystackChatMessage.from_user(translated_content)
                    )
                elif msg.role in ["bot", "assistant"]:
                    messages_for_llm.append(
                        HaystackChatMessage.from_assistant(translated_content)
                    )
                elif msg.role == "system":
                    messages_for_llm.append(
                        HaystackChatMessage.from_system(translated_content)
                    )
                else:
                    messages_for_llm.append(
                        HaystackChatMessage.from_user(translated_content)
                    )

            # --- Queue-based streaming generator ---
            def token_stream():
                collected_tokens = []
                token_queue = queue.Queue()
                rag_error = None

                def on_token(token):
                    collected_tokens.append(token)
                    token_queue.put(("token", token))

                # Run RAG in a separate thread
                def run_rag():
                    nonlocal rag_error
                    try:
                        get_rag_response_dynamic(
                            user_query_english=query,
                            chat_history_english=messages_for_llm,
                            on_token=on_token,
                        )
                    except Exception as e:
                        rag_error = e
                        logger.error(f"RAG processing error: {e}")
                    finally:
                        token_queue.put(("end", None))  # Signal end

                # Start RAG processing thread
                rag_thread = threading.Thread(target=run_rag)
                rag_thread.daemon = True  # Thread will die when main thread dies
                rag_thread.start()

                # Stream tokens as they arrive
                try:
                    while True:
                        try:
                            # Wait for token with timeout
                            message_type, token = token_queue.get(
                                timeout=60
                            )  # 60 second timeout

                            if message_type == "end":
                                break
                            elif message_type == "token" and token:
                                # Yield token in SSE format
                                yield f"data: {token}\n\n"

                        except queue.Empty:
                            logger.warning("Token queue timeout - no tokens received")
                            break

                    # Wait for RAG thread to complete (with timeout)
                    rag_thread.join(timeout=10)

                    # Handle any errors that occurred during RAG processing
                    if rag_error:
                        yield f"data: [ERROR] {str(rag_error)}\n\n"
                        return

                    # Save final response after streaming completes
                    if collected_tokens:
                        final_response = "".join(collected_tokens)

                        # Translate response if needed
                        translated_response = final_response
                        if language != "eng":
                            try:
                                translated_response = translate_text(
                                    "eng", language, final_response, sunbird_headers
                                )
                            except Exception as e:
                                logger.error(f"Translation error: {e}")
                                translated_response = (
                                    final_response  # Fallback to English
                                )

                        # Save bot message to database
                        ChatMessage.objects.create(
                            session=session,
                            role="bot",
                            content=translated_response,
                            language=language,
                        )
                        session.save(update_fields=["updated_at"])

                    # Send end marker
                    yield "data: [END]\n\n"

                except Exception as e:
                    logger.error(f"Streaming error: {e}")
                    yield f"data: [ERROR] Streaming failed: {str(e)}\n\n"

            # Return streaming response with appropriate headers
            response = StreamingHttpResponse(
                token_stream(), content_type="text/plain; charset=utf-8"
            )
            response["Cache-Control"] = "no-cache"
            response["Connection"] = "keep-alive"
            response["X-Accel-Buffering"] = "no"  # Disable nginx buffering
            return response

        except ValueError as e:
            logger.error(f"Chat processing error: {e}")
            return Response({"error": str(e)}, status=status.HTTP_400_BAD_REQUEST)
        except RuntimeError as e:
            logger.error(f"RAG service error: {e}")
            return Response(
                {"error": str(e)}, status=status.HTTP_503_SERVICE_UNAVAILABLE
            )
        except Exception as e:
            logger.exception("Unhandled error in ChatView:")
            return Response(
                {"error": "Internal server error."},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )


class KnowledgeBaseViewSet(viewsets.ModelViewSet):
    """
    API endpoint that allows knowledge base items to be viewed or uploaded.
    """

    queryset = KnowledgeBase.objects.all().order_by("-uploaded_at")
    serializer_class = KnowledgeBaseSerializer
    permission_classes = [IsAuthenticated]  # Uncomment and use later

    @action(detail=True, methods=["post"])
    def publish(self, request, pk=None):
        kb_instance = self.get_object()

        # Ensure only PDF is processed
        if not kb_instance.document or not kb_instance.document.name.lower().endswith(
            ".pdf"
        ):
            return Response(
                {"detail": "Only PDF documents are supported for publishing."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        if kb_instance.is_published:
            return Response(
                {"detail": "This file is already published."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        try:
            publish_to_neo4j(kb_instance)
            return Response(
                {"detail": "Publishing started successfully."},
                status=status.HTTP_200_OK,
            )
        except Exception as e:
            return Response(
                {"detail": f"Error while publishing: {str(e)}"},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )


class KnowledgeBaseStatsAPIView(APIView):
    def get(self, request, *args, **kwargs):
        stats = {
            "Documents": {"count": 0, "size_bytes": 0},
            "Images": {"count": 0, "size_bytes": 0},
            "Videos": {"count": 0, "size_bytes": 0},
            "Audio": {"count": 0, "size_bytes": 0},
        }

        for item in KnowledgeBase.objects.all():
            if item.document and os.path.exists(item.document.path):
                stats["Documents"]["count"] += 1
                stats["Documents"]["size_bytes"] += item.document.size
            if item.image and os.path.exists(item.image.path):
                stats["Images"]["count"] += 1
                stats["Images"]["size_bytes"] += item.image.size

            if item.video and os.path.exists(item.video.path):
                stats["Videos"]["count"] += 1
                stats["Videos"]["size_bytes"] += item.video.size

            if item.audio and os.path.exists(item.audio.path):
                stats["Audio"]["count"] += 1
                stats["Audio"]["size_bytes"] += item.audio.size

        def format_size(size_bytes):
            if size_bytes == 0:
                return "0 B"
            sizes = ["B", "KB", "MB", "GB", "TB"]
            i = 0
            while size_bytes >= 1024 and i < len(sizes) - 1:
                size_bytes /= 1024.0
                i += 1
            return f"{size_bytes:.2f} {sizes[i]}"

        # Format response for frontend
        data = []
        colors = {
            "Documents": "#3B82F6",
            "Images": "#00FF00",
            "Videos": "#F59E0B",
            "Audio": "#8B5CF6",
        }
        ALLOCATED_SPACE_GB = {
            "Documents": 10,  # 10GB allocated for documents
            "Images": 1,  # 1GB for images
            "Videos": 50,  # 50GB for videos
            "Audio": 5,  # 5GB for audio
        }

        for file_type, info in stats.items():
            allocated = ALLOCATED_SPACE_GB[file_type]
            size_gb = info["size_bytes"] / (1024**3)
            progress_percentage = (size_gb / allocated) * 100

            data.append(
                {
                    "name": file_type,
                    "size": format_size(info["size_bytes"]),
                    "files": info["count"],
                    "color": colors[file_type],
                    "progress": min(100, progress_percentage),
                    "allocated_space": f"{allocated} GB",
                }
            )
        # Add this before the return statement
        print("Stats data:", stats)
        for file_type, info in stats.items():
            print(
                f"{file_type}: count={info['count']}, size_bytes={info['size_bytes']}"
            )
        return Response(data, status=status.HTTP_200_OK)


class EscalateToAgentAPIView(APIView):
    def post(self, request, session_id):
        session = get_object_or_404(ChatSession, pk=session_id, user=request.user)
        session.state = "waiting"
        session.save(update_fields=["state"])
        # Notify agents via WebSocket (publish to /agent/queue)
        channel_layer = get_channel_layer()
        async_to_sync(channel_layer.group_send)(
            "agents",
            {"type": "new.request", "session_id": session.id, "title": session.title},
        )
        return Response({"detail": "Escalation request sent."})


class AgentAcceptAPIView(APIView):
    permission_classes = [IsAuthenticated]

    def post(self, request, session_id):
        session = get_object_or_404(ChatSession, pk=session_id, state="waiting")
        session.assigned_agent = request.user
        session.state = "agent"
        session.save(update_fields=["assigned_agent", "state"])
        # Notify citizen (group_send to session_{id})
        channel_layer = get_channel_layer()
        async_to_sync(channel_layer.group_send)(
            f"session_{session.id}",
            {"type": "agent.joined", "agent": request.user.username},
        )
        return Response({"detail": "You have joined the chat."})
