import logging
from uuid import UUID

from flask import request
from flask_restx import Resource
from pydantic import BaseModel, ConfigDict, Field, RootModel
from sqlalchemy.orm import Session
from werkzeug.datastructures import FileStorage
from werkzeug.exceptions import HTTPException, InternalServerError

import services
from controllers.common.fields import AudioTranscriptResponse
from controllers.common.rbac import AgentId, PlainApp, RBACCheck, enforce_rbac_checks
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.console import console_ns
from controllers.console.agent.app_helpers import (
    resolve_agent_runtime_app_model,
    resolve_existing_agent_runtime_app_model,
)
from controllers.console.app.error import (
    AppUnavailableError,
    AudioTooLargeError,
    CompletionRequestError,
    NoAudioUploadedError,
    ProviderModelCurrentlyNotSupportError,
    ProviderNotInitializeError,
    ProviderNotSupportSpeechToTextError,
    ProviderNotSupportTextToSpeechError,
    ProviderQuotaExceededError,
    SpeechToTextDisabledError,
    UnsupportedAudioTypeError,
)
from controllers.console.app.wraps import get_app_model, with_session
from controllers.console.wraps import (
    RBACPermission,
    account_initialization_required,
    edit_permission_required,
    model_validate,
    rbac_permission_required,
    setup_required,
    with_current_tenant_id,
    with_current_user,
)
from core.base.tts.audio_mime import SUPPORTED_TTS_AUDIO_MIME_TYPES
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from extensions.ext_database import db
from fields.base import ResponseModel
from graphon.model_runtime.errors.invoke import InvokeError
from libs.external_api import error_responses
from libs.flask_restx_compat import BINARY_RESPONSE_MEDIA_TYPES_VENDOR_KEY
from libs.helper import dump_response
from libs.login import current_user, login_required
from models import Account, App, AppMode
from models.agent import AgentConfigDraftType
from models.agent_config_entities import AgentSoulConfig
from services.agent.composer_service import AgentComposerService
from services.agent.errors import AgentNotFoundError, AgentVersionNotFoundError
from services.app_ref_service import AppRefService
from services.audio_service import AudioService
from services.errors.audio import (
    AudioTooLargeServiceError,
    NoAudioUploadedServiceError,
    ProviderNotSupportSpeechToTextServiceError,
    ProviderNotSupportTextToSpeechServiceError,
    SpeechToTextDisabledServiceError,
    UnsupportedAudioTypeServiceError,
)

logger = logging.getLogger(__name__)


class TextToSpeechPayload(BaseModel):
    message_id: str | None = Field(default=None, description="Message ID")
    text: str = Field(..., description="Text to convert")
    voice: str | None = Field(default=None, description="Voice name")
    streaming: bool | None = Field(default=None, description="Whether to stream audio")


class TextToSpeechVoiceQuery(BaseModel):
    language: str = Field(..., description="Language code")


class AgentAudioTranscriptFormPayload(BaseModel):
    model_config = ConfigDict(extra="forbid")

    draft_type: AgentConfigDraftType = AgentConfigDraftType.DRAFT


class TextToSpeechVoiceResponse(ResponseModel):
    # see api/core/plugin/impl/model.py
    name: str = Field(description="Voice display name")
    value: str = Field(description="Voice identifier")


class TextToSpeechVoiceListResponse(RootModel[list[TextToSpeechVoiceResponse]]):
    root: list[TextToSpeechVoiceResponse] = Field(description="Available voices")


register_schema_models(console_ns, AgentAudioTranscriptFormPayload, TextToSpeechPayload, TextToSpeechVoiceQuery)
register_response_schema_models(
    console_ns,
    AudioTranscriptResponse,
    TextToSpeechVoiceResponse,
    TextToSpeechVoiceListResponse,
)

_AUDIO_TRANSCRIPT_FILE_PARAM = {
    "description": "MP3 audio to transcribe",
    "in": "formData",
    "type": "file",
    "required": True,
}
_AGENT_AUDIO_TRANSCRIPT_PARAMS = {
    "file": _AUDIO_TRANSCRIPT_FILE_PARAM,
    "draft_type": {
        "description": "Agent debug config source",
        "in": "formData",
        "type": "string",
        "enum": [draft_type.value for draft_type in AgentConfigDraftType],
        "default": AgentConfigDraftType.DRAFT.value,
        "required": False,
    },
}
_CONSOLE_AUDIO_TRANSCRIPT_APP_MODES = [
    AppMode.CHAT,
    AppMode.AGENT_CHAT,
    AppMode.ADVANCED_CHAT,
    AppMode.AGENT,
]


def _transcribe_audio_to_text(
    *,
    app_model: App,
    file: FileStorage | None,
    session: Session,
    agent_soul: AgentSoulConfig | None = None,
) -> dict[str, str]:
    try:
        if agent_soul is None:
            response = AudioService.transcript_asr(
                app_model=app_model,
                file=file,
                session=session,
                end_user=None,
            )
        else:
            response = AudioService.transcript_agent_asr(
                app_model=app_model,
                agent_soul=agent_soul,
                file=file,
                session=session,
                end_user=None,
            )
        return dump_response(AudioTranscriptResponse, response)
    except services.errors.app_model_config.AppModelConfigBrokenError as caught_error:
        logger.exception("App model config broken.")
        raise AppUnavailableError() from caught_error
    except NoAudioUploadedServiceError as caught_error:
        raise NoAudioUploadedError() from caught_error
    except AudioTooLargeServiceError as e:
        raise AudioTooLargeError(str(e)) from e
    except UnsupportedAudioTypeServiceError as caught_error:
        raise UnsupportedAudioTypeError() from caught_error
    except ProviderNotSupportSpeechToTextServiceError as caught_error:
        raise ProviderNotSupportSpeechToTextError() from caught_error
    except SpeechToTextDisabledServiceError as caught_error:
        raise SpeechToTextDisabledError() from caught_error
    except ProviderTokenNotInitError as ex:
        raise ProviderNotInitializeError(ex.description) from ex
    except QuotaExceededError as caught_error:
        raise ProviderQuotaExceededError() from caught_error
    except ModelCurrentlyNotSupportError as caught_error:
        raise ProviderModelCurrentlyNotSupportError() from caught_error
    except InvokeError as e:
        raise CompletionRequestError(e.description) from e
    except HTTPException:
        raise
    except ValueError:
        raise
    except Exception as e:
        logger.exception("Failed to transcribe audio to text")
        raise InternalServerError() from e


@console_ns.route("/apps/<uuid:app_id>/audio-to-text")
class ChatMessageAudioApi(Resource):
    @error_responses(
        AgentVersionNotFoundError,
        AppUnavailableError,
        AudioTooLargeError,
        CompletionRequestError,
        NoAudioUploadedError,
        ProviderModelCurrentlyNotSupportError,
        ProviderNotInitializeError,
        ProviderNotSupportSpeechToTextError,
        ProviderQuotaExceededError,
        SpeechToTextDisabledError,
        UnsupportedAudioTypeError,
    )
    @console_ns.doc("chat_message_audio_transcript")
    @console_ns.doc(description="Transcript audio to text for chat messages")
    @console_ns.doc(
        consumes=["multipart/form-data"],
        params={"app_id": "App ID", "file": _AUDIO_TRANSCRIPT_FILE_PARAM},
    )
    @console_ns.response(
        200,
        "Audio transcription successful",
        console_ns.models[AudioTranscriptResponse.__name__],
    )
    @console_ns.response(400, "Bad request - No audio uploaded or unsupported type")
    @console_ns.response(413, "Audio file too large")
    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model(mode=_CONSOLE_AUDIO_TRANSCRIPT_APP_MODES)
    def post(self, app_model: App):
        return _transcribe_audio_to_text(
            app_model=app_model,
            file=request.files.get("file"),
            session=db.session(),
        )


@console_ns.route("/agent/<uuid:agent_id>/audio-to-text")
class AgentChatMessageAudioApi(Resource):
    @error_responses(
        AgentNotFoundError,
        AgentVersionNotFoundError,
        AppUnavailableError,
        AudioTooLargeError,
        CompletionRequestError,
        NoAudioUploadedError,
        ProviderModelCurrentlyNotSupportError,
        ProviderNotInitializeError,
        ProviderNotSupportSpeechToTextError,
        ProviderQuotaExceededError,
        SpeechToTextDisabledError,
        UnsupportedAudioTypeError,
    )
    @console_ns.doc("agent_chat_message_audio_transcript")
    @console_ns.doc(description="Transcribe audio using the current Agent debug configuration")
    @console_ns.doc(
        consumes=["multipart/form-data"],
        params={"agent_id": "Agent ID", **_AGENT_AUDIO_TRANSCRIPT_PARAMS},
    )
    @console_ns.response(
        200,
        "Audio transcription successful",
        console_ns.models[AudioTranscriptResponse.__name__],
    )
    @console_ns.response(400, "Bad request - Speech to text disabled or unsupported audio")
    @console_ns.response(404, "Agent or build draft not found")
    @console_ns.response(413, "Audio file too large")
    @setup_required
    @login_required
    @account_initialization_required
    @edit_permission_required
    @with_current_user
    @with_current_tenant_id
    @with_session
    def post(
        self,
        session: Session,
        current_tenant_id: str,
        current_user: Account,
        agent_id: UUID,
    ):
        payload = AgentAudioTranscriptFormPayload.model_validate(request.form.to_dict(flat=True))
        app_model = resolve_agent_runtime_app_model(
            session=session,
            tenant_id=current_tenant_id,
            agent_id=agent_id,
        )
        # Agent routes expose Agent ids, while APP RBAC is keyed by the resolved runtime App id.
        enforce_rbac_checks(
            tenant_id=current_tenant_id,
            account_id=current_user.id,
            checks=[RBACCheck(RBACPermission.APP_TEST_AND_RUN, PlainApp())],
            path_args={"app_id": app_model.id},
        )
        agent_soul = AgentComposerService.load_agent_soul_for_debug(
            tenant_id=current_tenant_id,
            agent_id=str(agent_id),
            account_id=current_user.id,
            draft_type=payload.draft_type,
            session=session,
        )
        return _transcribe_audio_to_text(
            app_model=app_model,
            agent_soul=agent_soul,
            file=request.files.get("file"),
            session=session,
        )


def _text_to_speech(*, payload: TextToSpeechPayload, app_model: App, session: Session, account_id: str | None):
    try:
        message_ref = None
        if payload.message_id:
            app_ref = AppRefService.create_app_ref(app_model)
            message_ref = AppRefService.create_message_ref(
                app_ref,
                payload.message_id,
                account_id=account_id,
            )

        return AudioService.transcript_tts(
            app_model=app_model,
            session=session,
            text=payload.text,
            voice=payload.voice,
            message_ref=message_ref,
            is_draft=True,
        )
    except services.errors.app_model_config.AppModelConfigBrokenError as caught_error:
        logger.exception("App model config broken.")
        raise AppUnavailableError() from caught_error
    except ProviderNotSupportTextToSpeechServiceError as caught_error:
        raise ProviderNotSupportTextToSpeechError() from caught_error
    except ProviderTokenNotInitError as ex:
        raise ProviderNotInitializeError(ex.description) from ex
    except QuotaExceededError as caught_error:
        raise ProviderQuotaExceededError() from caught_error
    except ModelCurrentlyNotSupportError as caught_error:
        raise ProviderModelCurrentlyNotSupportError() from caught_error
    except InvokeError as e:
        raise CompletionRequestError(e.description) from e
    except HTTPException:
        raise
    except ValueError:
        raise
    except Exception as e:
        logger.exception("Failed to convert text to speech")
        raise InternalServerError() from e


@console_ns.route("/apps/<uuid:app_id>/text-to-audio")
class ChatMessageTextApi(Resource):
    @error_responses(
        AppUnavailableError,
        CompletionRequestError,
        ProviderModelCurrentlyNotSupportError,
        ProviderNotInitializeError,
        ProviderNotSupportTextToSpeechError,
        ProviderQuotaExceededError,
    )
    @console_ns.doc("chat_message_text_to_speech")
    @console_ns.doc(description="Convert text to speech for chat messages")
    @console_ns.doc(params={"app_id": "App ID"})
    @console_ns.expect(console_ns.models[TextToSpeechPayload.__name__])
    # TTS returns provider audio bytes, so the success response is intentionally schema-less.
    @console_ns.response(200, "Text to speech conversion successful")
    @console_ns.response(400, "Bad request - Invalid parameters")
    @setup_required
    @login_required
    @account_initialization_required
    @get_app_model
    @model_validate(TextToSpeechPayload)
    def post(self, req_data: TextToSpeechPayload, app_model: App):
        # response-contract:ignore provider audio bytes
        return _text_to_speech(
            payload=req_data,
            app_model=app_model,
            session=db.session(),
            account_id=current_user.id if req_data.message_id else None,
        )


@console_ns.route("/agent/<uuid:agent_id>/text-to-audio")
class AgentChatMessageTextApi(Resource):
    @error_responses(
        AgentNotFoundError,
        AppUnavailableError,
        CompletionRequestError,
        ProviderModelCurrentlyNotSupportError,
        ProviderNotInitializeError,
        ProviderNotSupportTextToSpeechError,
        ProviderQuotaExceededError,
    )
    @console_ns.doc("agent_chat_message_text_to_speech")
    @console_ns.doc(description="Preview an Agent TTS voice or read an Agent chat message")
    @console_ns.doc(params={"agent_id": "Agent ID"})
    @console_ns.doc(
        produces=list(SUPPORTED_TTS_AUDIO_MIME_TYPES),
        vendor={BINARY_RESPONSE_MEDIA_TYPES_VENDOR_KEY: list(SUPPORTED_TTS_AUDIO_MIME_TYPES)},
    )
    @console_ns.expect(console_ns.models[TextToSpeechPayload.__name__])
    @console_ns.response(200, "Generated audio bytes in the provider audio format")
    @console_ns.response(400, "Invalid text or voice")
    @console_ns.response(403, "Insufficient permissions")
    @console_ns.response(404, "Agent not found")
    @setup_required
    @login_required
    @account_initialization_required
    @edit_permission_required
    @with_current_user
    @with_current_tenant_id
    @with_session(write=False)
    @model_validate(TextToSpeechPayload)
    def post(
        self,
        req_data: TextToSpeechPayload,
        session: Session,
        current_tenant_id: str,
        current_user: Account,
        agent_id: UUID,
    ):
        enforce_rbac_checks(
            tenant_id=current_tenant_id,
            account_id=current_user.id,
            checks=[RBACCheck(RBACPermission.AGENT_TEST_AND_RUN, AgentId())],
            path_args={"agent_id": str(agent_id)},
        )
        app_model = resolve_existing_agent_runtime_app_model(
            session=session,
            tenant_id=current_tenant_id,
            agent_id=agent_id,
        )
        # response-contract:ignore provider audio bytes
        return _text_to_speech(payload=req_data, app_model=app_model, session=session, account_id=current_user.id)


def _get_text_to_speech_voices(*, tenant_id: str, language: str):
    try:
        response = AudioService.transcript_tts_voices(
            tenant_id=tenant_id,
            language=language,
        )

        return dump_response(TextToSpeechVoiceListResponse, response)
    except ProviderTokenNotInitError as ex:
        raise ProviderNotInitializeError(ex.description) from ex
    except QuotaExceededError as caught_error:
        raise ProviderQuotaExceededError() from caught_error
    except ModelCurrentlyNotSupportError as caught_error:
        raise ProviderModelCurrentlyNotSupportError() from caught_error
    except InvokeError as e:
        raise CompletionRequestError(e.description) from e
    except ValueError as e:
        raise e
    except Exception as e:
        logger.exception("Failed to handle get request to TextModesApi")
        raise InternalServerError() from e


@console_ns.route("/apps/<uuid:app_id>/text-to-audio/voices")
class TextModesApi(Resource):
    @error_responses(
        CompletionRequestError,
        ProviderModelCurrentlyNotSupportError,
        ProviderNotInitializeError,
        ProviderQuotaExceededError,
    )
    @console_ns.doc("get_text_to_speech_voices")
    @console_ns.doc(description="Get available TTS voices for a specific language")
    @console_ns.doc(params={"app_id": "App ID", **query_params_from_model(TextToSpeechVoiceQuery)})
    @console_ns.response(
        200,
        "TTS voices retrieved successfully",
        console_ns.models[TextToSpeechVoiceListResponse.__name__],
    )
    @console_ns.response(400, "Invalid language parameter")
    @setup_required
    @login_required
    @account_initialization_required
    @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp()))
    @get_app_model
    @model_validate(TextToSpeechVoiceQuery)
    def get(self, req_data: TextToSpeechVoiceQuery, app_model: App):
        return _get_text_to_speech_voices(tenant_id=app_model.tenant_id, language=req_data.language)


@console_ns.route("/agent/<uuid:agent_id>/text-to-audio/voices")
class AgentTextToSpeechVoicesApi(Resource):
    @error_responses(
        AgentNotFoundError,
        CompletionRequestError,
        ProviderModelCurrentlyNotSupportError,
        ProviderNotInitializeError,
        ProviderQuotaExceededError,
    )
    @console_ns.doc("get_agent_text_to_speech_voices")
    @console_ns.doc(description="Get available TTS voices for an Agent and language")
    @console_ns.doc(params={"agent_id": "Agent ID", **query_params_from_model(TextToSpeechVoiceQuery)})
    @console_ns.response(
        200,
        "TTS voices retrieved successfully",
        console_ns.models[TextToSpeechVoiceListResponse.__name__],
    )
    @console_ns.response(400, "Invalid language parameter")
    @console_ns.response(404, "Agent not found")
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @with_session(write=False)
    @model_validate(TextToSpeechVoiceQuery)
    def get(
        self,
        req_data: TextToSpeechVoiceQuery,
        session: Session,
        current_tenant_id: str,
        current_user: Account,
        agent_id: UUID,
    ):
        enforce_rbac_checks(
            tenant_id=current_tenant_id,
            account_id=current_user.id,
            checks=[RBACCheck(RBACPermission.AGENT_PREVIEW, AgentId())],
            path_args={"agent_id": str(agent_id)},
        )
        app_model = resolve_existing_agent_runtime_app_model(
            session=session,
            tenant_id=current_tenant_id,
            agent_id=agent_id,
        )
        return _get_text_to_speech_voices(tenant_id=app_model.tenant_id, language=req_data.language)