"""Word答案解析节点：从.docx文件中提取题干和标准答案（带缓存）"""
import os
import json
import re
import logging
import tempfile
import requests
import orjson
from typing import List
from langchain_core.runnables import RunnableConfig
from langgraph.runtime import Runtime
from coze_coding_utils.runtime_ctx.context import Context
from coze_coding_dev_sdk import LLMClient
from langchain_core.messages import HumanMessage
from docx import Document

from graphs.state import (
    DocExtractInput,
    DocExtractOutput,
    CorrectAnswer
)
from utils.cache_manager import answer_doc_cache

logger = logging.getLogger(__name__)


def sanitize_json_string(text: str) -> str:
    """清理JSON字符串中的无效转义和控制字符"""
    text = ''.join(char if ord(char) >= 32 or char in '\n\r\t' else ' ' for char in text)
    result = []
    i = 0
    in_string = False
    while i < len(text):
        char = text[i]
        if char == '"' and (i == 0 or (i > 0 and result and result[-1] != '\\')):
            in_string = not in_string
            result.append(char)
        elif in_string and char == '\\':
            if i + 1 < len(text):
                next_char = text[i + 1]
                if next_char in ['"', '\\', '/', 'b', 'f', 'n', 'r', 't']:
                    result.append(char)
                    result.append(next_char)
                    i += 1
                elif next_char == 'u':
                    result.append(char)
                    result.append(next_char)
                    i += 1
                else:
                    result.append('\\\\')
            else:
                result.append('\\\\')
        elif in_string and char in '\n\r':
            result.append('\\n' if char == '\n' else '\\r')
        else:
            result.append(char)
        i += 1
    return ''.join(result)


def extract_json_from_text(text: str, key: str = "answers") -> dict:
    """从文本中提取JSON对象，多层回退策略"""
    
    # 先尝试直接解析完整JSON
    try:
        result = json.loads(text)
        if key in result:
            return result
    except (json.JSONDecodeError, ValueError) as e:
        logger.debug(f"Direct JSON parse failed: {e}")
    
    # 尝试找到JSON对象的边界
    try:
        # 找到第一个 { 和最后一个 }
        start = text.find('{')
        end = text.rfind('}')
        if start != -1 and end != -1 and end > start:
            json_str = text[start:end+1]
            result = json.loads(json_str)
            if key in result:
                return result
    except (json.JSONDecodeError, ValueError) as e:
        logger.debug(f"Boundary JSON parse failed: {e}")
    
    # 使用orjson尝试
    try:
        start = text.find('{')
        end = text.rfind('}')
        if start != -1 and end != -1 and end > start:
            json_str = text[start:end+1]
            result = orjson.loads(json_str)
            if key in result:
                return result
    except Exception as e:
        logger.debug(f"Orjson parse failed: {e}")
    
    # 尝试清理后解析
    try:
        cleaned = sanitize_json_string(text)
        start = cleaned.find('{')
        end = cleaned.rfind('}')
        if start != -1 and end != -1 and end > start:
            json_str = cleaned[start:end+1]
            result = json.loads(json_str)
            if key in result:
                return result
    except Exception as e:
        logger.debug(f"Cleaned JSON parse failed: {e}")
    
    logger.warning(f"Failed to extract JSON with key '{key}' from text (length: {len(text)})")
    return {key: []}


def download_and_extract_docx(url: str) -> str:
    """下载Word文件并提取文本内容"""
    logger.info(f"Downloading Word document from: {url}")
    
    tmp_path = None
    
    try:
        # 下载文件
        response = requests.get(url, timeout=60, allow_redirects=True)
        response.raise_for_status()
        
        # 检查内容类型
        content_type = response.headers.get('Content-Type', '')
        logger.debug(f"Response Content-Type: {content_type}")
        
        # 检查文件大小
        if len(response.content) < 100:
            raise ValueError(f"Downloaded file too small: {len(response.content)} bytes")
        
        # 检查是否为有效的 docx（ZIP 格式，以 PK 开头）
        if not response.content.startswith(b'PK'):
            # 可能是 HTML 错误页面
            content_preview = response.content[:1000].lower()
            if b'<html' in content_preview or b'<!doctype' in content_preview:
                raise ValueError("URL returned HTML instead of docx (may need authentication or URL expired)")
            raise ValueError("Downloaded file is not a valid docx (not ZIP format)")
        
        # 保存到临时文件
        with tempfile.NamedTemporaryFile(suffix='.docx', delete=False) as tmp_file:
            tmp_file.write(response.content)
            tmp_path = tmp_file.name
        
        # 使用python-docx解析
        doc = Document(tmp_path)
        
        # 提取所有段落文本
        paragraphs = []
        for para in doc.paragraphs:
            text = para.text.strip()
            if text:
                paragraphs.append(text)
        
        # 提取表格中的文本
        for table in doc.tables:
            for row in table.rows:
                row_text = []
                for cell in row.cells:
                    cell_text = cell.text.strip()
                    if cell_text:
                        row_text.append(cell_text)
                if row_text:
                    paragraphs.append(" | ".join(row_text))
        
        doc_text = "\n".join(paragraphs)
        logger.info(f"Extracted Word document text length: {len(doc_text)}")
        
        return doc_text
        
    except Exception as e:
        logger.error(f"Failed to download/extract docx: {e}")
        raise
    finally:
        # 清理临时文件
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except Exception:
                pass


def parse_answer_doc_with_llm(answer_doc_url: str, ctx, config: RunnableConfig) -> List[CorrectAnswer]:
    """
    使用LLM解析答案文档（实际解析逻辑）
    
    Args:
        answer_doc_url: 答案文档URL
        ctx: 上下文
        config: 配置
        
    Returns:
        解析后的正确答案列表
    """
    # 1. 下载并提取Word文档内容
    try:
        doc_text = download_and_extract_docx(answer_doc_url)
    except Exception as e:
        logger.error(f"Failed to download/extract Word document: {e}")
        return []
    
    if not doc_text.strip():
        logger.error("No text content extracted from Word document")
        return []
    
    logger.info(f"Word document content preview: {doc_text[:500]}")
    
    # 2. 使用LLM解析题目和答案
    cfg_file = os.path.join(os.getenv("COZE_WORKSPACE_PATH", ""), config["metadata"]["llm_cfg"])
    with open(cfg_file, "r", encoding="utf-8") as fd:
        _cfg = json.load(fd)
    
    llm_config = _cfg.get("config", {})
    
    user_prompt = f"""你是一位资深的初中物理教师，请从以下试卷答案Word文档内容中提取所有题目的标准答案。

【Word文档内容】
{doc_text[:20000]}

【提取要求】
1. 识别所有题号，包括大题和小题
2. 只提取每道题的标准答案，不需要提取题干
3. 答案格式：
   - 选择题：单个字母（A/B/C/D）
   - 填空题：数值或表达式
   - 解答题：关键结果
4. 如果有多个空，用逗号分隔

【题号格式规范】
- 大题题号：直接使用数字，如"1"、"2"、"10"
- 小题题号：使用"大题.小题"格式，如"10.1"、"10.2"

【重要】
- 保持答案简洁，不要包含题干
- 如果文档中有分值，请提取；否则默认3分

请返回简洁的JSON格式：
{{
  "answers": [
    {{
      "question_id": "1",
      "correct_answer": "B",
      "full_score": 3
    }},
    {{
      "question_id": "2",
      "correct_answer": "2a-b+c",
      "full_score": 3
    }}
  ]
}}"""

    client = LLMClient(ctx=ctx)
    response = client.invoke(
        messages=[HumanMessage(content=user_prompt)],
        model=llm_config.get("model", "doubao-seed-2-0-pro-260215"),
        temperature=llm_config.get("temperature", 0.1),
        max_completion_tokens=llm_config.get("max_completion_tokens", 8192)
    )
    
    response_text = response.content if isinstance(response.content, str) else " ".join(
        item.get("text", "") if isinstance(item, dict) else str(item) 
        for item in response.content
    ).strip()
    
    logger.info(f"Word extract LLM response: {response_text[:1500]}")
    
    # 清理markdown标记
    for prefix in ["```json", "```JSON", "```"]:
        if response_text.startswith(prefix):
            response_text = response_text[len(prefix):]
    for suffix in ["```"]:
        if response_text.endswith(suffix):
            response_text = response_text[:-3]
    
    # 解析JSON
    result_dict = extract_json_from_text(response_text.strip(), "answers")
    
    correct_answers: List[CorrectAnswer] = []
    for ans in result_dict.get("answers", []):
        try:
            correct_answers.append(CorrectAnswer(
                question_id=str(ans.get("question_id", "")),
                parent_id=str(ans.get("parent_id", "")),
                is_sub_question=bool(ans.get("is_sub_question", False)),
                question_text=str(ans.get("question_text", "")),
                correct_answer=str(ans.get("correct_answer", "")),
                full_score=int(ans.get("full_score", 3) if ans.get("full_score") is not None else 3),
                answer_analysis=str(ans.get("answer_analysis", ""))
            ))
        except Exception as e:
            logger.warning(f"Failed to parse correct answer: {ans}, error: {e}")
            continue
    
    logger.info(f"Parsed {len(correct_answers)} correct answers from Word document")
    
    # 打印解析结果供调试
    for ans in correct_answers:
        logger.info(f"  Question {ans.question_id}: {ans.correct_answer} ({ans.full_score}分)")
    
    return correct_answers


def doc_extract_node(
    state: DocExtractInput,
    config: RunnableConfig,
    runtime: Runtime[Context]
) -> DocExtractOutput:
    """
    title: Word答案解析
    desc: 从正确答案Word文件（.docx）中提取题干和标准答案，用于后续批改；如果未提供URL则返回空列表；支持缓存，避免重复解析
    integrations: 大语言模型
    """
    ctx = runtime.context
    
    # 检查是否提供了答案文档URL
    if not state.answer_doc_url or not state.answer_doc_url.strip():
        logger.info("No answer document URL provided, will use teacher mode for all questions")
        return DocExtractOutput(correct_answers=[])
    
    # 尝试从缓存获取
    cached_result = answer_doc_cache.get(state.answer_doc_url)
    if cached_result is not None:
        logger.info(f"Cache hit for answer doc: {state.answer_doc_url[:50]}...")
        # 将字典列表转换回CorrectAnswer对象
        correct_answers = [
            CorrectAnswer(**ans) if isinstance(ans, dict) else ans
            for ans in cached_result
        ]
        return DocExtractOutput(correct_answers=correct_answers)
    
    logger.info(f"Cache miss for answer doc: {state.answer_doc_url[:50]}...")
    
    # 缓存未命中，执行解析
    correct_answers = parse_answer_doc_with_llm(state.answer_doc_url, ctx, config)
    
    # 存入缓存（将CorrectAnswer对象转换为字典列表）
    if correct_answers:
        cache_data = [
            {
                "question_id": ans.question_id,
                "parent_id": ans.parent_id,
                "is_sub_question": ans.is_sub_question,
                "question_text": ans.question_text,
                "correct_answer": ans.correct_answer,
                "full_score": ans.full_score,
                "answer_analysis": ans.answer_analysis
            }
            for ans in correct_answers
        ]
        answer_doc_cache.set(state.answer_doc_url, cache_data)
        logger.info(f"Cached {len(correct_answers)} answers for: {state.answer_doc_url[:50]}...")
    
    return DocExtractOutput(correct_answers=correct_answers)
