# eveai_api/api/specialist_execution_api.py
import uuid

from flask import Response, stream_with_context, current_app
from flask_restx import Namespace, Resource, fields
from flask_jwt_extended import jwt_required, get_jwt_identity

from common.utils.celery_utils import current_celery
from common.utils.execution_progress import ExecutionProgressTracker
from eveai_api.api.auth import requires_service

specialist_execution_ns = Namespace('specialist-execution', description='Specialist execution operations')

specialist_start_session_response = specialist_execution_ns.model('StartSessionResponse', {
    'session_id': fields.String(required=True, description='A new Chat session ID'),
})


@specialist_execution_ns.route('/start_session', methods=['GET'])
class StartSession(Resource):
    @jwt_required()
    @requires_service("SPECIALIST_API")
    @specialist_execution_ns.response(201, 'New Session ID created Successfully', specialist_start_session_response)
    def get(self):
        new_session_id = f"{uuid.uuid4()}"
        return {
            'session_id': new_session_id,
        }, 201


specialist_execution_input = specialist_execution_ns.model('SpecialistExecutionInput', {
    'specialist_id': fields.Integer(required=True, description='ID of the specialist to use'),
    'arguments': fields.Raw(required=True, description='Dynamic arguments for specialist and retrievers'),
    'session_id': fields.String(required=True, description='Chat session ID'),
    'user_timezone': fields.String(required=True, description='User timezone')
})

specialist_execution_response = specialist_execution_ns.model('SpecialistExecutionResponse', {
    'task_id': fields.String(description='ID of specialist execution task, to be used to retrieve execution stream'),
    'status': fields.String(description='Status of the execution'),
    'stream_url': fields.String(description='Stream URL'),
})


@specialist_execution_ns.route('')
class StartExecution(Resource):
    @jwt_required()
    @requires_service('SPECIALIST_API')
    @specialist_execution_ns.expect(specialist_execution_input)
    @specialist_execution_ns.response(201, 'Specialist execution successfully queued.', specialist_execution_response)
    def post(self):
        """Start execution of a specialist"""
        tenant_id = get_jwt_identity()
        data = specialist_execution_ns.payload

        # Send task to queue
        task = current_celery.send_task(
            'execute_specialist',
            args=[tenant_id,
                  data['specialist_id'],
                  data['arguments'],
                  data['session_id'],
                  data['user_timezone'],
                  ],
            queue='llm_interactions'
        )

        return {
            'task_id': task.id,
            'status': 'queued',
            'stream_url': f'/api/v1/specialist-execution/{task.id}/stream'
        }, 201


@specialist_execution_ns.route('/<string:task_id>/stream')
class ExecutionStream(Resource):
    @jwt_required()
    @requires_service('SPECIALIST_API')
    def get(self, task_id: str):
        """Get streaming updates for a specialist execution"""
        progress_tracker = ExecutionProgressTracker()
        return Response(
            stream_with_context(progress_tracker.get_updates(task_id)),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'Connection': 'keep-alive'
            }
        )