"""
LLM 信息提取

使用 LLM 从 PDF 文本中提取结构化数据，以及从支付截图中提取刷卡信息。
支持 JSON 格式输出，字段与 CSV_COLUMNS 对齐。

对外接口:
    extract_invoice_from_text(text, file_name) -> dict    从 PDF 文本提取发票信息
    extract_card_info_from_image(image_path) -> dict     从支付截图提取刷卡信息
"""

from __future__ import annotations

import base64
import json
from pathlib import Path
from typing import Any, cast

from .. import get_logger
from .prompt import build_card_info_system_prompt, build_invoice_system_prompt

log = get_logger("llm_extractor")


def _create_llm() -> Any:
    """根据配置文件创建 LLM 实例。"""
    try:
        from llama_index.llms.openai_like import OpenAILike
    except ImportError:
        log.error("缺少 llama-index-llms-openai-like，请执行: uv pip install llama-index-llms-openai-like")
        raise

    from ..config import get_llm_config

    llm_config = get_llm_config()
    return OpenAILike(
        model=llm_config["model"],
        api_base=llm_config["api_base"],
        api_key=llm_config.get("api_key", "lm-studio"),
        temperature=0.1,
        max_tokens=8192,
        request_timeout=600.0,
        is_chat_model=True,
    )


def _llm_query(system_prompt: str, user_content: str, max_tokens: int = 4096) -> str:
    """发送请求到 LLM 并返回完整响应文本。"""
    from llama_index.core.llms import ChatMessage

    from ..config import get_llm_config

    messages = [
        ChatMessage(role="system", content=system_prompt),
        ChatMessage(role="user", content=user_content),
    ]

    llm_config = get_llm_config()
    llm = _create_llm()
    log.info("开始请求 LLM (model=%s, base=%s)", llm_config["model"], llm_config["api_base"])

    try:
        parts = []
        for resp in llm.stream_chat(
            messages,
            temperature=0.1,
            max_tokens=max_tokens,
            extra_body={"reasoning_effort": "none"},
        ):
            delta = resp.delta
            if delta:
                parts.append(delta)
        text = "".join(parts)
        log.info("LLM 请求完成，响应总长度: %d 字符", len(text))
        log.info("LLM 响应: %s", text)
        return text
    except Exception as e:
        log.error("LLM 请求失败: %s", e)
        raise


def _parse_json_response(text: str) -> dict[str, Any]:
    """从 LLM 响应中提取 JSON，处理可能的 Markdown 包裹。"""
    text = text.strip()

    # 处理 ```json ... ``` 包裹
    if "```" in text:
        # 提取第一个代码块
        start = text.find("```") + 3
        end = text.find("```", start)
        if end > start:
            text = text[start:end].strip()

    # 去掉可能的前缀 (如 "json")
    if text.lower().startswith("json"):
        text = text[4:].strip()

    return cast(dict[str, Any], json.loads(text))


def extract_invoice_from_text(text: str, file_name: str = "") -> dict[str, Any]:
    """从 PDF 发票文本中提取结构化数据。

    Args:
        text: PDF 提取的文本内容。
        file_name: 原始文件名（用于日志）。

    Returns:
        包含所有 CSV_COLUMNS 字段的字典。
    """
    system_prompt = build_invoice_system_prompt()
    user_content = f"请分析以下发票文本并提取信息：\n\n文件名: {file_name}\n\n---\n\n{text}\n\n---"

    try:
        response = _llm_query(system_prompt, user_content, max_tokens=4096)
        result = _parse_json_response(response)
        log.info("LLM 发票提取成功: %s", file_name)
        return result
    except Exception as e:
        log.error("LLM 发票提取失败: %s (%s)", file_name, e)
        raise


# ------------------------------------------------------------------
# 支付截图信息提取（多模态）
# ------------------------------------------------------------------


def _image_to_base64(image_path: Path) -> str:
    """将图片文件读取为 base64 字符串。"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def _llm_query_multimodal(
    system_prompt: str,
    text: str,
    image_b64: str,
    max_tokens: int = 4096,
) -> str:
    """发送多模态请求（文本 + 图片）到 LLM。"""
    from llama_index.core.base.llms.types import ImageBlock, TextBlock
    from llama_index.core.llms import ChatMessage

    from ..config import get_llm_config

    messages = [
        ChatMessage(role="system", content=system_prompt),
        ChatMessage(
            role="user",
            blocks=[
                TextBlock(text=text),
                ImageBlock(
                    url=f"data:image/jpeg;base64,{image_b64}",
                    detail="high",
                ),
            ],
        ),
    ]

    llm_config = get_llm_config()
    llm = _create_llm()
    log.info(
        "开始请求 LLM 多模态 (model=%s, base=%s)",
        llm_config["model"],
        llm_config["api_base"],
    )

    try:
        parts = []
        for resp in llm.stream_chat(
            messages,
            temperature=0.1,
            max_tokens=max_tokens,
            extra_body={"reasoning_effort": "none"},
        ):
            delta = resp.delta
            if delta:
                parts.append(delta)
        text = "".join(parts)
        log.info("LLM 多模态请求完成，响应总长度: %d 字符", len(text))
        log.info("LLM 多模态响应: %s", text)
        return text
    except Exception as e:
        log.error("LLM 多模态请求失败: %s", e)
        raise


def extract_card_info_from_image(image_path: Path) -> dict[str, Any]:
    """从支付截图中提取刷卡信息。

    Args:
        image_path: 支付截图图片路径。

    Returns:
        包含刷卡日期、刷卡金额、公务卡号的字典。
    """
    system_prompt = build_card_info_system_prompt()
    user_text = f"请分析以下支付截图并提取信息：\n\n文件名: {image_path.name}"

    image_b64 = _image_to_base64(image_path)

    try:
        response = _llm_query_multimodal(system_prompt, user_text, image_b64, max_tokens=4096)
        result = _parse_json_response(response)
        log.info("LLM 支付截图提取成功: %s", image_path.name)
        return result
    except Exception as e:
        log.error("LLM 支付截图提取失败: %s (%s)", image_path.name, e)
        raise
