from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from typing import Any

import httpx
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.models.ai_config import ModelConfig
from app.services.external_errors import ExternalServiceError
from app.services.model_service import (
    _anthropic_headers,
    _auth_headers,
    _call_configured_model,
    _copy_model_with_overrides,
    _decimal_to_float,
    _load_extra_params,
    _mock_answer,
    _max_output_tokens,
    _openai_messages,
    _put_if_not_none,
    _resolve_anthropic_endpoint,
    _resolve_openai_endpoint,
    _rough_token_count,
    _system_and_turn_messages,
    _system_config_bool,
)
from app.services.model_routing_service import ModelRoute, ModelRoutingService
from app.services.rag_service import RagResult


@dataclass
class StreamingModelResponse:
    model_id: int | None
    model_name: str
    input_token: int
    chunks: Iterator[str]
    route_reason: str | None = None
    question_type: str | None = None


@dataclass
class AsyncStreamingModelResponse:
    model_id: int | None
    model_name: str
    input_token: int
    chunks: AsyncIterator[str]
    route_reason: str | None = None
    question_type: str | None = None


class ModelStreamService:
    @staticmethod
    def stream(db: Session, rag_result: RagResult) -> StreamingModelResponse:
        route = _chat_route(db, rag_result)
        model = route.model
        mock_model_enabled = _system_config_bool(db, "mock_model_enabled", get_settings().mock_model_enabled)

        if mock_model_enabled:
            model_name = model.model_name if model is not None else "mock-model"
            return StreamingModelResponse(
                model_id=model.id if model is not None else None,
                model_name=model_name,
                input_token=_rough_token_count(rag_result.prompt),
                chunks=_display_chunks(model, _mock_answer(rag_result)),
                route_reason=route.reason,
                question_type=route.question_type,
            )

        model, fallback, route_reason = _prepare_routed_model(db, route)
        response = StreamingModelResponse(
            model_id=model.id,
            model_name=model.model_name,
            input_token=_rough_token_count(rag_result.prompt),
            chunks=iter(()),
            route_reason=route_reason,
            question_type=route.question_type,
        )
        response.chunks = _stream_with_runtime_fallback(response, model, fallback, rag_result)
        return response

    @staticmethod
    def stream_async(db: Session, rag_result: RagResult) -> AsyncStreamingModelResponse:
        route = _chat_route(db, rag_result)
        model = route.model
        mock_model_enabled = _system_config_bool(db, "mock_model_enabled", get_settings().mock_model_enabled)

        if mock_model_enabled:
            model_name = model.model_name if model is not None else "mock-model"
            return AsyncStreamingModelResponse(
                model_id=model.id if model is not None else None,
                model_name=model_name,
                input_token=_rough_token_count(rag_result.prompt),
                chunks=_async_display_chunks(model, _mock_answer(rag_result)),
                route_reason=route.reason,
                question_type=route.question_type,
            )

        model, fallback, route_reason = _prepare_routed_model(db, route)
        response = AsyncStreamingModelResponse(
            model_id=model.id,
            model_name=model.model_name,
            input_token=_rough_token_count(rag_result.prompt),
            chunks=_empty_async_chunks(),
            route_reason=route_reason,
            question_type=route.question_type,
        )
        response.chunks = _stream_with_runtime_fallback_async(response, model, fallback, rag_result)
        return response

    @staticmethod
    def debug_stream_async(
        model: ModelConfig,
        rag_result: RagResult,
        overrides: dict[str, Any],
        *,
        fallback_model: ModelConfig | None = None,
        route_reason: str = "后台调试：管理员手动指定模型",
        question_type: str | None = None,
    ) -> AsyncStreamingModelResponse:
        debug_model = _copy_model_with_overrides(model, overrides)
        debug_fallback = (
            _copy_model_with_overrides(fallback_model, overrides)
            if fallback_model is not None and fallback_model.id != model.id
            else None
        )
        if not _is_configured(debug_model):
            if debug_fallback is None or not _is_configured(debug_fallback):
                raise ExternalServiceError("模型 Base URL/API URL 或 API Key 未配置", provider="model")
            debug_model = debug_fallback
            route_reason = f"{route_reason}；场景模型配置不可用，运行时回退默认主模型"
        response = AsyncStreamingModelResponse(
            model_id=debug_model.id,
            model_name=debug_model.model_name,
            input_token=_rough_token_count(rag_result.prompt),
            chunks=_empty_async_chunks(),
            route_reason=route_reason,
            question_type=question_type,
        )
        response.chunks = _stream_with_runtime_fallback_async(
            response,
            debug_model,
            debug_fallback,
            rag_result,
        )
        return response


def _chat_route(db: Session, rag_result: RagResult) -> ModelRoute:
    return ModelRoutingService.resolve_chat(
        db,
        [chunk.knowledge_type for chunk in rag_result.chunks],
        rag_result.question,
    )


def _prepare_routed_model(
    db: Session,
    route: ModelRoute,
) -> tuple[ModelConfig, ModelConfig | None, str]:
    model = route.model
    if model is None:
        raise ExternalServiceError("未启用可用模型，请先在模型管理中启用一个模型。", provider="model")
    fallback = _runtime_fallback_model(db, model)
    if _is_configured(model):
        return model, fallback, route.reason
    if fallback is None or not _is_configured(fallback):
        raise ExternalServiceError("模型 Base URL/API URL 或 API Key 未配置", provider="model")
    return fallback, fallback, f"{route.reason}；场景模型配置不可用，运行时回退默认主模型"


def _runtime_fallback_model(db: Session, selected: ModelConfig) -> ModelConfig | None:
    default = ModelRoutingService.default_model(db)
    if default is None or default.id == selected.id:
        return None
    return default


def _is_configured(model: ModelConfig) -> bool:
    return bool((model.api_url or model.base_url) and model.api_key)


def _stream_with_runtime_fallback(
    response: StreamingModelResponse,
    model: ModelConfig,
    fallback: ModelConfig | None,
    rag_result: RagResult,
) -> Iterator[str]:
    emitted = False
    try:
        for chunk in _stream_configured_model(model, rag_result):
            if chunk:
                emitted = True
            yield chunk
        return
    except ExternalServiceError:
        if emitted or fallback is None or fallback.id == model.id or not _is_configured(fallback):
            raise
    response.model_id = fallback.id
    response.model_name = fallback.model_name
    response.route_reason = f"{response.route_reason}；场景模型调用失败，运行时回退默认主模型"
    yield from _stream_configured_model(fallback, rag_result)


async def _stream_with_runtime_fallback_async(
    response: AsyncStreamingModelResponse,
    model: ModelConfig,
    fallback: ModelConfig | None,
    rag_result: RagResult,
) -> AsyncIterator[str]:
    emitted = False
    try:
        async for chunk in _stream_configured_model_async(model, rag_result):
            if chunk:
                emitted = True
            yield chunk
        return
    except ExternalServiceError:
        if emitted or fallback is None or fallback.id == model.id or not _is_configured(fallback):
            raise
    response.model_id = fallback.id
    response.model_name = fallback.model_name
    response.route_reason = f"{response.route_reason}；场景模型调用失败，运行时回退默认主模型"
    async for chunk in _stream_configured_model_async(fallback, rag_result):
        yield chunk


async def _empty_async_chunks() -> AsyncIterator[str]:
    if False:
        yield ""


def _stream_configured_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]:
    api_type = model.api_type or "openai_compatible"
    if model.stream_enabled != 1:
        return iter((_call_configured_model(model, rag_result, allow_no_hit=True),))
    if api_type == "anthropic_messages":
        return _stream_anthropic_messages(model, rag_result)
    if api_type == "openai_compatible":
        return _stream_openai_compatible_model(model, rag_result)
    return iter((_call_configured_model(model, rag_result, allow_no_hit=True),))


async def _stream_configured_model_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]:
    api_type = model.api_type or "openai_compatible"
    if model.stream_enabled != 1:
        answer = await asyncio.to_thread(
            _call_configured_model, model, rag_result, allow_no_hit=True
        )
        yield answer
        return

    if api_type == "anthropic_messages":
        async for chunk in _stream_anthropic_messages_async(model, rag_result):
            yield chunk
        return

    if api_type == "openai_compatible":
        async for chunk in _stream_openai_compatible_model_async(model, rag_result):
            yield chunk
        return

    answer = await asyncio.to_thread(
        _call_configured_model, model, rag_result, allow_no_hit=True
    )
    yield answer


def _stream_openai_compatible_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]:
    payload: dict[str, Any] = {
        "model": model.model_name,
        "messages": _openai_messages(rag_result, model),
        "max_tokens": _max_output_tokens(model),
    }
    _put_if_not_none(payload, "temperature", _decimal_to_float(model.temperature))
    _put_if_not_none(payload, "top_p", _decimal_to_float(model.top_p))
    _put_if_not_none(payload, "presence_penalty", _decimal_to_float(model.presence_penalty))
    _put_if_not_none(payload, "frequency_penalty", _decimal_to_float(model.frequency_penalty))
    if model.response_format == "json_object":
        payload["response_format"] = {"type": "json_object"}
    payload.update(_load_extra_params(model.extra_params))
    payload["stream"] = True

    try:
        with httpx.stream(
            "POST",
            _resolve_openai_endpoint(model),
            json=payload,
            headers=_auth_headers(model),
            timeout=model.timeout_second,
        ) as response:
            response.raise_for_status()
            yield from _iter_openai_stream(response)
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
        raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc


async def _stream_openai_compatible_model_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]:
    payload = _openai_stream_payload(model, rag_result)
    try:
        async with httpx.AsyncClient(timeout=model.timeout_second) as client:
            async with client.stream(
                "POST",
                _resolve_openai_endpoint(model),
                json=payload,
                headers=_auth_headers(model),
            ) as response:
                response.raise_for_status()
                async for chunk in _iter_openai_stream_async(response):
                    yield chunk
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
        raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc


def _openai_stream_payload(model: ModelConfig, rag_result: RagResult) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "model": model.model_name,
        "messages": _openai_messages(rag_result, model),
        "max_tokens": _max_output_tokens(model),
    }
    _put_if_not_none(payload, "temperature", _decimal_to_float(model.temperature))
    _put_if_not_none(payload, "top_p", _decimal_to_float(model.top_p))
    _put_if_not_none(payload, "presence_penalty", _decimal_to_float(model.presence_penalty))
    _put_if_not_none(payload, "frequency_penalty", _decimal_to_float(model.frequency_penalty))
    if model.response_format == "json_object":
        payload["response_format"] = {"type": "json_object"}
    payload.update(_load_extra_params(model.extra_params))
    payload["stream"] = True
    return payload


def _iter_openai_stream(response: httpx.Response) -> Iterator[str]:
    in_reasoning = False
    for data_line in _iter_sse_data(response):
        data = json.loads(data_line)
        if isinstance(data, dict) and data.get("error"):
            raise ValueError(str(data["error"]))

        choices = data.get("choices")
        if not isinstance(choices, list) or not choices:
            continue
        first_choice = choices[0]
        if not isinstance(first_choice, dict):
            continue

        delta = first_choice.get("delta")
        if isinstance(delta, dict):
            reasoning = _string_or_empty(delta.get("reasoning_content") or delta.get("reasoning"))
            if reasoning:
                if not in_reasoning:
                    in_reasoning = True
                    yield "<think>"
                yield reasoning
                continue

            content = _string_or_empty(delta.get("content"))
            if content:
                if in_reasoning:
                    in_reasoning = False
                    yield "</think>"
                yield content
                continue

        text = _string_or_empty(first_choice.get("text"))
        if text:
            if in_reasoning:
                in_reasoning = False
                yield "</think>"
            yield text

    if in_reasoning:
        yield "</think>"


async def _iter_openai_stream_async(response: httpx.Response) -> AsyncIterator[str]:
    in_reasoning = False
    async for data_line in _iter_sse_data_async(response):
        data = json.loads(data_line)
        if isinstance(data, dict) and data.get("error"):
            raise ValueError(str(data["error"]))

        choices = data.get("choices")
        if not isinstance(choices, list) or not choices:
            continue
        first_choice = choices[0]
        if not isinstance(first_choice, dict):
            continue

        delta = first_choice.get("delta")
        if isinstance(delta, dict):
            reasoning = _string_or_empty(delta.get("reasoning_content") or delta.get("reasoning"))
            if reasoning:
                if not in_reasoning:
                    in_reasoning = True
                    yield "<think>"
                yield reasoning
                continue

            content = _string_or_empty(delta.get("content"))
            if content:
                if in_reasoning:
                    in_reasoning = False
                    yield "</think>"
                yield content
                continue

        text = _string_or_empty(first_choice.get("text"))
        if text:
            if in_reasoning:
                in_reasoning = False
                yield "</think>"
            yield text

    if in_reasoning:
        yield "</think>"


def _stream_anthropic_messages(model: ModelConfig, rag_result: RagResult) -> Iterator[str]:
    system, messages = _system_and_turn_messages(rag_result, model)
    payload: dict[str, Any] = {
        "model": model.model_name,
        "max_tokens": _max_output_tokens(model),
        "system": system,
        "messages": messages,
    }
    _put_if_not_none(payload, "temperature", _decimal_to_float(model.temperature))
    _put_if_not_none(payload, "top_p", _decimal_to_float(model.top_p))
    _put_if_not_none(payload, "top_k", model.top_k)
    payload.update(_load_extra_params(model.extra_params))
    payload["stream"] = True

    try:
        with httpx.stream(
            "POST",
            _resolve_anthropic_endpoint(model),
            json=payload,
            headers=_anthropic_headers(model),
            timeout=model.timeout_second,
        ) as response:
            response.raise_for_status()
            yield from _iter_anthropic_stream(response)
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
        raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc


async def _stream_anthropic_messages_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]:
    payload = _anthropic_stream_payload(model, rag_result)
    try:
        async with httpx.AsyncClient(timeout=model.timeout_second) as client:
            async with client.stream(
                "POST",
                _resolve_anthropic_endpoint(model),
                json=payload,
                headers=_anthropic_headers(model),
            ) as response:
                response.raise_for_status()
                async for chunk in _iter_anthropic_stream_async(response):
                    yield chunk
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
        raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc


def _anthropic_stream_payload(model: ModelConfig, rag_result: RagResult) -> dict[str, Any]:
    system, messages = _system_and_turn_messages(rag_result, model)
    payload: dict[str, Any] = {
        "model": model.model_name,
        "max_tokens": _max_output_tokens(model),
        "system": system,
        "messages": messages,
    }
    _put_if_not_none(payload, "temperature", _decimal_to_float(model.temperature))
    _put_if_not_none(payload, "top_p", _decimal_to_float(model.top_p))
    _put_if_not_none(payload, "top_k", model.top_k)
    payload.update(_load_extra_params(model.extra_params))
    payload["stream"] = True
    return payload


def _iter_anthropic_stream(response: httpx.Response) -> Iterator[str]:
    in_reasoning = False
    for data_line in _iter_sse_data(response):
        data = json.loads(data_line)
        if isinstance(data, dict) and data.get("type") == "error":
            raise ValueError(str(data.get("error", data)))

        if not isinstance(data, dict) or data.get("type") != "content_block_delta":
            continue
        delta = data.get("delta")
        if not isinstance(delta, dict):
            continue

        if delta.get("type") == "thinking_delta":
            thinking = _string_or_empty(delta.get("thinking"))
            if thinking:
                if not in_reasoning:
                    in_reasoning = True
                    yield "<think>"
                yield thinking
            continue

        text = _string_or_empty(delta.get("text"))
        if text:
            if in_reasoning:
                in_reasoning = False
                yield "</think>"
            yield text

    if in_reasoning:
        yield "</think>"


async def _iter_anthropic_stream_async(response: httpx.Response) -> AsyncIterator[str]:
    in_reasoning = False
    async for data_line in _iter_sse_data_async(response):
        data = json.loads(data_line)
        if isinstance(data, dict) and data.get("type") == "error":
            raise ValueError(str(data.get("error", data)))

        if not isinstance(data, dict) or data.get("type") != "content_block_delta":
            continue
        delta = data.get("delta")
        if not isinstance(delta, dict):
            continue

        if delta.get("type") == "thinking_delta":
            thinking = _string_or_empty(delta.get("thinking"))
            if thinking:
                if not in_reasoning:
                    in_reasoning = True
                    yield "<think>"
                yield thinking
            continue

        text = _string_or_empty(delta.get("text"))
        if text:
            if in_reasoning:
                in_reasoning = False
                yield "</think>"
            yield text

    if in_reasoning:
        yield "</think>"


def _iter_sse_data(response: httpx.Response) -> Iterator[str]:
    for line in response.iter_lines():
        if not line:
            continue
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        line = line.strip()
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if data == "[DONE]":
            break
        if data:
            yield data


async def _iter_sse_data_async(response: httpx.Response) -> AsyncIterator[str]:
    async for line in response.aiter_lines():
        if not line:
            continue
        line = line.strip()
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if data == "[DONE]":
            break
        if data:
            yield data


def _chunk_text(text: str, *, chunk_size: int = 12) -> Iterator[str]:
    for index in range(0, len(text), chunk_size):
        yield text[index : index + chunk_size]


async def _async_chunk_text(text: str, *, chunk_size: int = 12) -> AsyncIterator[str]:
    for index in range(0, len(text), chunk_size):
        yield text[index : index + chunk_size]
        await asyncio.sleep(0)


def _display_chunks(model: ModelConfig | None, text: str) -> Iterator[str]:
    return _chunk_text(text) if model is None or model.stream_enabled == 1 else iter((text,))


async def _async_display_chunks(model: ModelConfig | None, text: str) -> AsyncIterator[str]:
    if model is not None and model.stream_enabled != 1:
        yield text
        return
    async for chunk in _async_chunk_text(text):
        yield chunk


def _string_or_empty(value: Any) -> str:
    return value if isinstance(value, str) else ""
