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 import select
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.rag_service import RagResult


@dataclass(frozen=True)
class StreamingModelResponse:
    model_id: int | None
    model_name: str
    input_token: int
    chunks: Iterator[str]


@dataclass(frozen=True)
class AsyncStreamingModelResponse:
    model_id: int | None
    model_name: str
    input_token: int
    chunks: AsyncIterator[str]


class ModelStreamService:
    @staticmethod
    def stream(db: Session, rag_result: RagResult) -> StreamingModelResponse:
        model = _get_enabled_model(db)
        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)),
            )

        if model is None:
            raise ExternalServiceError("未启用可用模型，请先在模型管理中启用一个模型。", provider="model")
        if not (model.api_url or model.base_url) or not model.api_key:
            raise ExternalServiceError("模型 Base URL/API URL 或 API Key 未配置", provider="model")

        return StreamingModelResponse(
            model_id=model.id,
            model_name=model.model_name,
            input_token=_rough_token_count(rag_result.prompt),
            chunks=_stream_configured_model(model, rag_result),
        )

    @staticmethod
    def stream_async(db: Session, rag_result: RagResult) -> AsyncStreamingModelResponse:
        model = _get_enabled_model(db)
        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)),
            )

        if model is None:
            raise ExternalServiceError("未启用可用模型，请先在模型管理中启用一个模型。", provider="model")
        if not (model.api_url or model.base_url) or not model.api_key:
            raise ExternalServiceError("模型 Base URL/API URL 或 API Key 未配置", provider="model")

        return AsyncStreamingModelResponse(
            model_id=model.id,
            model_name=model.model_name,
            input_token=_rough_token_count(rag_result.prompt),
            chunks=_stream_configured_model_async(model, rag_result),
        )

    @staticmethod
    def debug_stream_async(
        model: ModelConfig,
        rag_result: RagResult,
        overrides: dict[str, Any],
    ) -> AsyncStreamingModelResponse:
        debug_model = _copy_model_with_overrides(model, overrides)
        if not (debug_model.api_url or debug_model.base_url) or not debug_model.api_key:
            raise ExternalServiceError("模型 Base URL/API URL 或 API Key 未配置", provider="model")
        return AsyncStreamingModelResponse(
            model_id=model.id,
            model_name=model.model_name,
            input_token=_rough_token_count(rag_result.prompt),
            chunks=_stream_configured_model_async(debug_model, rag_result),
        )


def _get_enabled_model(db: Session) -> ModelConfig | None:
    return db.scalar(
        select(ModelConfig)
        .where(ModelConfig.enabled == 1)
        .order_by(ModelConfig.id.desc())
        .limit(1)
    )


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 ""
