import json
from threading import Lock
import uuid
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

from flask import current_app , url_for
from flask_babel import _, force_locale  

from extensions import db
from Copilot.LLMLayer.Services.PromtBuilder import PromptBuilder
from Copilot.LLMLayer.Services.RetrievalService import RetrievalService
from Copilot.LLMLayer.Services.uiAnalyzerShocket import uiAnalyzerShocket
from Copilot.LLMLayer.Services.incidentDatabaseService import IncidentDatabaseService
from Copilot.dataAcquisition.Core.BaseSubscriber import BaseSubscriber
from Copilot.SituationLayer.Core.Incident import IncidentRef
from Copilot.LLMLayer.Core.LLMAnalysisRequest import LLMAnalysisRequest
from Copilot.LLMLayer.Utils.LLMUtils import incident_to_llm_context, llm_json_dumps



from services.chat_service import ChatService
from models.ai_model import AITaskModelConfig, ModelType


class LLMAnalyzerRequestHandler(BaseSubscriber):
    def __init__(self, socketio, app=None):
        super().__init__()
        self._lock = Lock()
        self.socketio = socketio
        self._executor = ThreadPoolExecutor(max_workers=8)
        
        if app is None:
            app = current_app._get_current_object()
        self.app = app
    
    def handle(self, event: LLMAnalysisRequest):
        self._executor.submit(self._process, event)

    def _process(self, event: LLMAnalysisRequest):
        socket_service = None
        incident = None
        try:
            user_msg = event.user_message
            incident: IncidentRef = event.incident
            user_id = event.user_id
            lang = event.lang

            incident_context_dict = incident_to_llm_context(incident)
            incident_context = llm_json_dumps(incident_context_dict)

            prompt_builder = PromptBuilder()
            system_prompt = prompt_builder.build_system_prompt(language=lang)
            incident_prompt = prompt_builder.build_incident_prompt(incident_context, user_msg)

            with self.app.app_context():
                with force_locale(lang):
                    MAX_STAGE = 4
                    socket_service = uiAnalyzerShocket(self.socketio)

                    text_model = AITaskModelConfig.get_model_for_task(ModelType.TEXT)
                    model_name = text_model.name

                    retrieval = RetrievalService(current_app.doc_manager, text_model_name="gpt-5-mini")

                    source_prompt = ""
                    query = ""

                    if retrieval.has_any_dcoument():
                        socket_service.emit(True, incident.id,
                                            stage_name=_("خلاصه سازی رخداد"), stage_number=1,
                                            max_stage_number=MAX_STAGE, description=_("ایجاد کوئری برای جست و جو"))
                        query = retrieval.build_query(incident_context_dict, user_msg)

                        socket_service.emit(True, incident.id,
                                            stage_name=_("بررسی اسناد"), stage_number=2,
                                            max_stage_number=MAX_STAGE, description=_("در حال جست و جو در اسناد"))
                        source_prompt = retrieval.retrieve(query)

                    socket_service.emit(True, incident.id,
                                        stage_name=_("بررسی رخداد"), stage_number=3,
                                        max_stage_number=MAX_STAGE, description=_("در حال بررسی نهایی و ارایه تحلیل"))

                    reply, ok = ChatService.ask_openai(
                        incident_prompt, history=[], resources=source_prompt,
                        text_model_name=model_name, system_prompt=system_prompt,
                        enable_web_search=text_model.supports_web_search,
                    )

                    if not ok:
                        socket_service.emit(False, incident.id, description=_('خطایی رخ داد'))
                        return

                    reply_json: dict = json.loads(reply)
                    title = reply_json.get('title', "new incident")

                    db_service = IncidentDatabaseService()
                    analysis = db_service.create_analysis(incident.id, reply_json, incident_context, user_msg, query, len(source_prompt))
                    session = db_service.create_chat_session(user_id, analysis.id, title)
                    redirect_url = f"/chat/{session.session_url}"
                    db.session.commit()

                    socket_service.emit(True, incident.id, redirect_to=redirect_url,
                                        stage_name=_("اتمام"), stage_number=MAX_STAGE,
                                        max_stage_number=MAX_STAGE, description=_("اتمام عملیات"))

        except Exception as e:
            if socket_service and incident:
                socket_service.emit(False, incident.id, description=str(e))