import logging
import logging.config
import os

import numpy as np

from lightrag import LightRAG
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
from openai import OpenAI

from Config.Config import *


async def print_stream(stream):
    async for chunk in stream:
        if chunk:
            print(chunk, end="", flush=True)


def configure_logging():
    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
        logger_instance = logging.getLogger(logger_name)
        logger_instance.handlers = []
        logger_instance.filters = []

    log_dir = os.getenv("LOG_DIR", os.getcwd())
    log_file_path = os.path.abspath(
        os.path.join(log_dir, "./Logs/lightrag.log")
    )

    print(f"\nLightRAG log file: {log_file_path}\n")
    os.makedirs(os.path.dirname(log_dir), exist_ok=True)

    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))
    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))

    logging.config.dictConfig(
        {
            "version": 1,
            "disable_existing_loggers": False,
            "formatters": {
                "default": {
                    "format": "%(levelname)s: %(message)s",
                },
                "detailed": {
                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                },
            },
            "handlers": {
                "console": {
                    "formatter": "default",
                    "class": "logging.StreamHandler",
                    "stream": "ext://sys.stderr",
                },
                "file": {
                    "formatter": "detailed",
                    "class": "logging.handlers.RotatingFileHandler",
                    "filename": log_file_path,
                    "maxBytes": log_max_bytes,
                    "backupCount": log_backup_count,
                    "encoding": "utf-8",
                },
            },
            "loggers": {
                "lightrag": {
                    "handlers": ["console", "file"],
                    "level": "INFO",
                    "propagate": False,
                },
            },
        }
    )
    logger.setLevel(logging.INFO)
    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")


async def llm_model_func(
        prompt, system_prompt=None, history_messages=None, **kwargs
) -> str:
    return await openai_complete_if_cache(
        os.getenv("LLM_MODEL", LLM_MODEL_NAME),
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        api_key=LLM_API_KEY,
        base_url=LLM_BASE_URL,
        **kwargs,
    )


async def embedding_func(texts: list[str]) -> np.ndarray:
    return await openai_embed(
        texts,
        model=EMBED_MODEL_NAME,
        api_key=EMBED_API_KEY,
        base_url=EMBED_BASE_URL
    )


async def initialize_rag(working_dir):
    rag = LightRAG(
        working_dir=working_dir,
        llm_model_func=llm_model_func,
        embedding_func=EmbeddingFunc(
            embedding_dim=EMBED_DIM,
            max_token_size=EMBED_MAX_TOKEN_SIZE,
            func=embedding_func
        ),
    )

    await rag.initialize_storages()
    await initialize_pipeline_status()

    return rag


def create_llm_model_func():
    def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
        return openai_complete_if_cache(
            LLM_MODEL_NAME,
            prompt,
            system_prompt=system_prompt,
            history_messages=history_messages,
            api_key=LLM_API_KEY,
            base_url=LLM_BASE_URL,
            **kwargs,
        )

    return llm_model_func


def create_embedding_func():
    return EmbeddingFunc(
        embedding_dim=1024,
        max_token_size=8192,
        func=lambda texts: openai_embed(
            texts,
            model=EMBED_MODEL_NAME,
            api_key=EMBED_API_KEY,
            base_url=EMBED_BASE_URL,
        ),
    )


def create_vision_model_func(llm_model_func):
    def vision_model_func(
            prompt, system_prompt=None, history_messages=[], image_data=None, **kwargs
    ):
        if image_data:
            return openai_complete_if_cache(
                VISION_MODEL_NAME,
                "",
                system_prompt=None,
                history_messages=[],
                messages=[
                    {"role": "system", "content": system_prompt}
                    if system_prompt
                    else None,
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_data}"
                                },
                            },
                        ],
                    }
                    if image_data
                    else {"role": "user", "content": prompt},
                ],
                api_key=VISION_API_KEY,
                base_url=VISION_BASE_URL,
                **kwargs,
            )
        else:
            return llm_model_func(prompt, system_prompt, history_messages, **kwargs)

    return vision_model_func


def format_exam_content(raw_text, output_path):
    client = OpenAI(
        api_key=LLM_API_KEY,
        base_url=LLM_BASE_URL
    )
    """
    将OCR识别的原始试卷内容格式化为标准试题格式
    
    参数:
        client: OpenAI客户端实例
        raw_text: OCR识别的原始文本
        output_path: 输出文件路径
    
    返回:
        格式化后的试题内容
    """
    prompt = """
我将提供一份markdown格式的试卷，请帮我整理出每道题的以下内容：
1. 题目序号
2. 题目内容（自动识别并添加$或$$包裹数学公式）
3. 选项（如果有）
4. 答案
5. 解析

要求：
- 一道题一道题输出，不要使用表格
- 自动检测数学表达式并用$或$$正确包裹
- 确保公式中的特殊字符正确转义
- 除题目内容外，不要输出其它无关信息

内容如下：
"""
    prompt += raw_text
    
    completion = client.chat.completions.create(
        model="deepseek-v3",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt},
        ],
    )
    
    formatted_content = completion.choices[0].message.content
    
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(formatted_content)