"""一体化识别批改节点：合并识别和批改为一次LLM调用，提升速度"""
import os
import json
import re
import logging
import orjson
from typing import List, Dict, Any
from jinja2 import Template
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 graphs.state import (
    RecognizeAndCorrectInput,
    RecognizeAndCorrectOutput,
    QuestionItem,
    CorrectionResult,
    MarkPosition,
    CorrectAnswer
)

logger = logging.getLogger(__name__)

# 思考过程的特征词
THINKING_KEYWORDS = [
    "不对", "重新看", "可能我", "哦，", "说明我", "这明显",
    "不，", "应该是", "我发现", "等等", "让我", "我理解",
    "？不", "？不对", "不是", "搞错了", "看错了"
]


def clean_comment(text: str) -> str:
    """清理comment中的思考过程"""
    for keyword in THINKING_KEYWORDS:
        if keyword in text:
            # 找到思考过程开始的位置
            idx = text.find(keyword)
            if idx > 0:
                # 截断思考过程
                cleaned = text[:idx]
                # 尝试保留完整句子（最后一个句号之后）
                last_period = cleaned.rfind("。")
                if last_period > 0:
                    return cleaned[:last_period + 1]
                return cleaned
    return text


def fix_incomplete_json(text: str) -> str:
    """尝试修复不完整的JSON字符串"""
    # 计算括号数量
    brace_count = text.count('{') - text.count('}')
    bracket_count = text.count('[') - text.count(']')
    
    # 补全缺失的括号
    if bracket_count > 0:
        text += ']' * bracket_count
    if brace_count > 0:
        text += '}' * brace_count
    
    return text


def extract_complete_objects(text: str) -> List[dict]:
    """从JSON文本中提取完整的对象列表（处理思考过程干扰）"""
    objects = []
    
    # 找到每个对象的开始位置
    obj_pattern = r'\{\s*"question_id"'
    matches = list(re.finditer(obj_pattern, text))
    
    for i, match in enumerate(matches):
        start = match.start()
        
        # 找到这个对象的结束位置
        brace_count = 0
        in_string = False
        escape = False
        end = start
        
        for j in range(start, len(text)):
            char = text[j]
            
            if escape:
                escape = False
                continue
            
            if char == '\\':
                escape = True
                continue
            
            if char == '"' and not escape:
                in_string = not in_string
                continue
            
            if in_string:
                continue
            
            if char == '{':
                brace_count += 1
            elif char == '}':
                brace_count -= 1
                if brace_count == 0:
                    end = j + 1
                    break
        
        if end > start:
            obj_str = text[start:end]
            try:
                obj = orjson.loads(obj_str)
                if isinstance(obj, dict) and "question_id" in obj:
                    objects.append(obj)
                    logger.debug(f"Extracted object: {obj.get('question_id')}")
            except Exception as e:
                logger.debug(f"Failed to parse object: {e}")
    
    return objects


def extract_json_from_text(text: str, key: str = "results") -> dict:
    """从文本中提取JSON对象，增强健壮性"""
    
    # 清理markdown标记
    for prefix in ["```json", "```JSON", "```"]:
        if text.startswith(prefix):
            text = text[len(prefix):]
    for suffix in ["```"]:
        if text.endswith(suffix):
            text = text[:-3]
    
    text = text.strip()
    
    # 尝试直接解析
    for parser in [json.loads, lambda x: orjson.loads(x)]:
        try:
            result = parser(text)
            if isinstance(result, dict) and key in result:
                logger.info("Successfully parsed JSON directly")
                return result
        except Exception as e:
            logger.debug(f"Direct parse failed: {e}")
    
    # 尝试修复不完整的JSON
    try:
        fixed_text = fix_incomplete_json(text)
        for parser in [json.loads, lambda x: orjson.loads(x)]:
            try:
                result = parser(fixed_text)
                if isinstance(result, dict) and key in result:
                    logger.info("Successfully parsed fixed JSON")
                    return result
            except:
                pass
    except:
        pass
    
    # 尝试提取完整的对象列表（处理思考过程干扰）
    objects = extract_complete_objects(text)
    if objects:
        logger.info(f"Extracted {len(objects)} complete objects from JSON with thinking")
        return {key: objects}
    
    logger.warning(f"Failed to extract JSON with key '{key}' from text length {len(text)}")
    return {key: []}


def build_dynamic_prompt(
    image_width: int,
    image_height: int,
    correct_answers: List[CorrectAnswer],
    comment_max_length: int
) -> str:
    """构建动态Prompt部分"""
    
    if correct_answers:
        answers_text = json.dumps(
            [{"question_id": a.question_id, "correct_answer": a.correct_answer} 
             for a in correct_answers],
            ensure_ascii=False
        )
        answer_hint = f"""
【标准答案】
{answers_text}"""
    else:
        answer_hint = "\n【批改模式】无标准答案，请根据物理知识判断。"
    
    return f"""
【图片尺寸】{image_width}×{image_height}像素
【评语限制】{comment_max_length}字以内
{answer_hint}

⚠️ 必须输出完整JSON，不要输出思考过程"""


def recognize_and_correct_node(
    state: RecognizeAndCorrectInput,
    config: RunnableConfig,
    runtime: Runtime[Context]
) -> RecognizeAndCorrectOutput:
    """
    title: 一体化识别批改
    desc: 合并识别和批改为一次LLM调用，提升批改速度
    integrations: 大语言模型
    """
    ctx = runtime.context
    
    # 读取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", {})
    sp = _cfg.get("sp", "")
    up = _cfg.get("up", "")
    
    # 获取参数
    image_url = state.image_url
    image_info = state.image_info
    correct_answers = state.correct_answers
    comment_max_length = getattr(state, 'comment_max_length', 100)
    
    # 使用Jinja2渲染up
    up_tpl = Template(up)
    user_prompt = up_tpl.render({
        "image_url": image_url,
        "comment_max_length": comment_max_length
    })
    
    # 构建动态部分
    dynamic_prompt = build_dynamic_prompt(
        image_info.width,
        image_info.height,
        correct_answers,
        comment_max_length
    )
    
    # 组合完整Prompt
    full_prompt = f"{sp}\n\n{user_prompt}\n{dynamic_prompt}"
    
    # 调用LLM
    messages = [
        HumanMessage(content=[
            {"type": "text", "text": full_prompt},
            {"type": "image_url", "image_url": {"url": image_url}}
        ])
    ]
    
    client = LLMClient(ctx=ctx)
    response = client.invoke(
        messages=messages,
        model=llm_config.get("model", "doubao-seed-2-0-pro-260215"),
        temperature=llm_config.get("temperature", 0.0),
        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()
    
    # 只记录前500字符，避免日志过大
    logger.info(f"LLM response (first 500 chars): {response_text[:500]}")
    
    # 解析结果
    result_dict = extract_json_from_text(response_text, "results")
    
    # 相对坐标（0-1000）转换为绝对坐标
    # 严格检查图片尺寸，防止除零和无效坐标
    img_width = max(1, image_info.width) if image_info.width > 0 else 1000
    img_height = max(1, image_info.height) if image_info.height > 0 else 1000
    
    width_scale = img_width / 1000.0
    height_scale = img_height / 1000.0
    
    question_items: List[QuestionItem] = []
    correction_results: List[CorrectionResult] = []
    
    results_list = result_dict.get("results", [])
    if not isinstance(results_list, list):
        results_list = []
    
    for r in results_list:
        try:
            if not isinstance(r, dict):
                continue
            
            # 解析bbox（相对坐标0-1000）
            answer_bbox = r.get("answer_bbox", [0, 0, 0, 0])
            if not isinstance(answer_bbox, list) or len(answer_bbox) != 4:
                answer_bbox = [0, 0, 0, 0]
            
            # 确保bbox值在有效范围内（0-1000）
            answer_bbox = [
                max(0, min(1000, int(answer_bbox[0]))) if isinstance(answer_bbox[0], (int, float)) else 0,
                max(0, min(1000, int(answer_bbox[1]))) if isinstance(answer_bbox[1], (int, float)) else 0,
                max(0, min(1000, int(answer_bbox[2]))) if isinstance(answer_bbox[2], (int, float)) else 0,
                max(0, min(1000, int(answer_bbox[3]))) if isinstance(answer_bbox[3], (int, float)) else 0
            ]
            
            # 确保x2 >= x1, y2 >= y1
            if answer_bbox[2] < answer_bbox[0]:
                answer_bbox[2] = answer_bbox[0]
            if answer_bbox[3] < answer_bbox[1]:
                answer_bbox[3] = answer_bbox[1]
            
            # 转换为绝对坐标
            answer_bbox_abs = [
                int(answer_bbox[0] * width_scale),
                int(answer_bbox[1] * height_scale),
                int(answer_bbox[2] * width_scale),
                int(answer_bbox[3] * height_scale)
            ]
            
            # 严格限制绝对坐标在图片范围内
            answer_bbox_abs = [
                max(0, min(img_width, answer_bbox_abs[0])),
                max(0, min(img_height, answer_bbox_abs[1])),
                max(0, min(img_width, answer_bbox_abs[2])),
                max(0, min(img_height, answer_bbox_abs[3]))
            ]
            
            # 自动计算mark_position（智能定位，紧贴答案框）
            bbox_width = answer_bbox_abs[2] - answer_bbox_abs[0]
            bbox_height = answer_bbox_abs[3] - answer_bbox_abs[1]
            
            if bbox_width > 0 and bbox_height > 0:
                # 计算答案框右侧的剩余空间
                right_space = img_width - answer_bbox_abs[2]
                
                # 策略1：如果右侧空间充足（>80px），标记紧贴答案框右侧
                if right_space > 80:
                    mark_x = answer_bbox_abs[2] + 10
                    mark_y = answer_bbox_abs[1] + int(bbox_height * 0.5)
                # 策略2：如果右侧空间不足但有一定空间（>40px），标记在答案框右上角内部
                elif right_space > 40:
                    mark_x = answer_bbox_abs[2] - 10
                    mark_y = answer_bbox_abs[1] + 10
                # 策略3：如果右侧空间很小，标记在答案框左上角
                else:
                    mark_x = answer_bbox_abs[0] + 10
                    mark_y = answer_bbox_abs[1] + 10
                
                # 最终边界检查（严格限制在图片范围内，留10px边距）
                mark_x = max(10, min(mark_x, img_width - 10))
                mark_y = max(10, min(mark_y, img_height - 10))
            else:
                # bbox无效时，使用图片中心
                mark_x = img_width // 2
                mark_y = img_height // 2
            
            mark_position = MarkPosition(x=mark_x, y=mark_y)
            
            # 构建QuestionItem
            question_items.append(QuestionItem(
                question_id=str(r.get("question_id", "")),
                parent_id=str(r.get("parent_id", "")),
                is_sub_question=bool(r.get("is_sub_question", False)),
                question_text=str(r.get("question_text", "")),
                student_answer=str(r.get("student_answer", "")),
                answer_bbox=answer_bbox_abs,
                mark_position=mark_position,
                full_score=int(r.get("full_score", 10) if r.get("full_score") is not None else 10)
            ))
            
            # 构建CorrectionResult
            status = str(r.get("status", "incorrect"))
            if status not in ["correct", "incorrect", "partial"]:
                status = "incorrect"
            
            # 清理comment中的思考过程
            comment_raw = str(r.get("comment", ""))
            comment = clean_comment(comment_raw)
            
            # 如果清理后comment为空，设置默认值
            if not comment:
                comment = "正确" if status == "correct" else "错误"
            
            correction_results.append(CorrectionResult(
                question_id=str(r.get("question_id", "")),
                parent_id=str(r.get("parent_id", "")),
                status=status,
                score=int(r.get("score", 0) if r.get("score") is not None else 0),
                full_score=int(r.get("full_score", 10) if r.get("full_score") is not None else 10),
                comment=comment
            ))
            
        except Exception as e:
            logger.warning(f"Failed to parse result: {r}, error: {e}")
            continue
    
    logger.info(f"Parsed {len(question_items)} questions and {len(correction_results)} correction results")
    
    return RecognizeAndCorrectOutput(
        question_items=question_items,
        correction_results=correction_results
    )
