import re
import threading
from collections import Counter

import gradio as gr
import spaces
import transformers
from transformers import pipeline

# 모델과 토크나이저 로딩
model_name = "CohereForAI/c4ai-command-r7b-arabic-02-2025"
if gr.NO_RELOAD:
    pipe = pipeline(
        "text-generation",
        model=model_name,
        device_map="auto",
        torch_dtype="auto",
    )

# 최종 답변을 감지하기 위한 마커
ANSWER_MARKER = "**답변**"

# 단계별 추론을 시작하는 문장들
rethink_prepends = [
    "자, 이제 다음을 파악해야 합니다 ",
    "제 생각에는 ",
    "잠시만요, 제 생각에는 ",
    "다음 사항이 맞는지 확인해 보겠습니다 ",
    "또한 기억해야 할 것은 ",
    "또 다른 주목할 점은 ",
    "그리고 저는 다음과 같은 사실도 기억합니다 ",
    "이제 충분히 이해했다고 생각합니다 ",
]

# 일반적인 추론 가이드 프롬프트
general_reasoning_guide = """
이 문제를 해결하기 위한 체계적인 접근 방법을 사용하겠습니다:

1. 문제에서 제공된 모든 정보와 조건을 명확히 이해합니다.
2. 각 변수와 관계를 식별하고 필요한 방정식을 세웁니다.
3. 단계별로 계산을 수행하며, 각 단계의 결과를 확인합니다.
4. 중간 결과가 합리적인지 검토하며 진행합니다.
5. 최종 답안을 도출하고 문제의 요구사항을 충족하는지 확인합니다.

이제 문제를 풀어보겠습니다:
"""

# 결과 추출 및 검증을 위한 함수들
def extract_calculation_results(reasoning_text):
    """추론 과정에서 도출된 가능한 답안 결과를 추출합니다."""
    # 수치 결과 패턴 (다양한 표현 방식 고려)
    numeric_patterns = [
        r'결과는 (\d+[\.,]?\d*)',
        r'답(은|는|이) (\d+[\.,]?\d*)',
        r'정답(은|는|이) (\d+[\.,]?\d*)',
        r'답안(은|는|이) (\d+[\.,]?\d*)',
        r'수익(은|는|이) (\d+[\.,]?\d*)',
        r'값(은|는|이) (\d+[\.,]?\d*)',
        r'결론(은|는|이) (\d+[\.,]?\d*)',
        r'개수(는|은|가) (\d+[\.,]?\d*)',
        r'총 (\d+[\.,]?\d*)개',
        r'총액(은|는|이) (\d+[\.,]?\d*)',
        r'총합(은|는|이) (\d+[\.,]?\d*)',
        r'합계(는|은|가) (\d+[\.,]?\d*)',
        r'=\s*(\d+[\.,]?\d*)\s*$',
        r':\s*(\d+[\.,]?\d*)\s*$',
        r'총계:\s*(\d+[\.,]?\d*)',
        r'최종 결과:\s*(\d+[\.,]?\d*)',
        r'최종 값:\s*(\d+[\.,]?\d*)',
        r'최종 답변:\s*(\d+[\.,]?\d*)',
    ]
    
    # 단위를 포함한 패턴 (달러, 개, 세트 등)
    unit_patterns = [
        r'(\d+[\.,]?\d*)\s*(달러|원|유로|파운드|엔)',
        r'(\d+[\.,]?\d*)\s*(개|명|세트|쌍|팀|그룹)',
        r'(\d+[\.,]?\d*)\s*(분|시간|초|일|주|개월|년)',
        r'(\d+[\.,]?\d*)\s*(미터|킬로미터|센티미터|인치|피트)',
        r'(\d+[\.,]?\d*)\s*(그램|킬로그램|파운드|온스)',
    ]
    
    results = []
    
    # 숫자 결과 추출
    for pattern in numeric_patterns:
        matches = re.findall(pattern, reasoning_text, re.IGNORECASE)
        for match in matches:
            if isinstance(match, tuple):
                # 그룹이 여러 개인 경우 (첫 번째는 조사 등)
                value = match[-1]  # 마지막 그룹이 숫자값
            else:
                value = match
            # 콤마 제거 및 소수점 처리
            value = value.replace(',', '')
            try:
                if '.' in value:
                    results.append(float(value))
                else:
                    results.append(int(value))
            except ValueError:
                continue
    
    # 단위가 포함된 결과 추출
    for pattern in unit_patterns:
        matches = re.findall(pattern, reasoning_text, re.IGNORECASE)
        for match in matches:
            value = match[0].replace(',', '')
            try:
                if '.' in value:
                    results.append(float(value))
                else:
                    results.append(int(value))
            except ValueError:
                continue
    
    # 마지막 문단에서 숫자만 추출 (최종 답변에 가까운 숫자)
    last_paragraph = reasoning_text.split('\n\n')[-1]
    numbers_in_last = re.findall(r'(\d+[\.,]?\d*)', last_paragraph)
    for num in numbers_in_last:
        num = num.replace(',', '')
        try:
            if '.' in num:
                results.append(float(num))
            else:
                results.append(int(num))
        except ValueError:
            continue
    
    return results

def determine_best_result(results, full_reasoning):
    """가장 신뢰할 수 있는 결과를 결정합니다."""
    if not results:
        return None
    
    # 결과가 하나밖에 없으면 그것을 반환
    if len(set(results)) == 1:
        return results[0]
    
    # 빈도 기반 분석 (가장 자주 등장한 결과가 신뢰성이 높을 가능성)
    counter = Counter(results)
    most_common = counter.most_common()
    
    # 빈도가 높은 상위 결과들
    top_results = [result for result, count in most_common if count >= most_common[0][1] * 0.8]
    
    if len(top_results) == 1:
        return top_results[0]
    
    # 최종 결론 근처에 있는 결과에 더 높은 가중치 부여
    paragraphs = full_reasoning.split('\n\n')
    last_paragraphs = '\n\n'.join(paragraphs[-2:])  # 마지막 두 단락
    
    # 마지막 단락에서 등장하는 결과 확인
    final_results = [result for result in top_results if str(result) in last_paragraphs]
    if final_results:
        # 마지막 단락에서 가장 자주 등장한 결과
        final_counter = Counter([r for r in results if r in final_results])
        if final_counter:
            return final_counter.most_common(1)[0][0]
    
    # 수식과 함께 등장하는 결과 (예: "= 78", "총합: 78")
    for result in top_results:
        result_str = str(result)
        if re.search(r'=\s*' + result_str + r'(?!\d)', full_reasoning) or \
           re.search(r'결과[:는은이가]\s*' + result_str, full_reasoning) or \
           re.search(r'답[:는은이가]\s*' + result_str, full_reasoning) or \
           re.search(r'정답[:는은이가]\s*' + result_str, full_reasoning):
            return result
    
    # 위의 방법으로 결정할 수 없을 경우 가장 빈도가 높은 결과 반환
    return most_common[0][0]

# 중간 결과를 요약하기 위한 프롬프트
structured_reasoning_prompt = """
지금까지의 추론을 단계별로 정리해보겠습니다:

1. 문제 분석:
   - 주어진 정보: {given_info}
   - 구해야 할 것: {goal}

2. 계산 과정:
{calculation_steps}

3. 현재까지의 결론:
{current_conclusion}

이제 다음 단계로 진행하겠습니다.
"""

# 최종 결과 검증을 위한 프롬프트
verification_prompt = """
지금까지의 추론 과정에서 여러 결과가 도출되었습니다:
{different_results}

이 중에서 가장 정확한 답변을 찾기 위해 계산 과정을 처음부터 다시 검토하겠습니다:

1. 문제 분석:
   - 주어진 정보: {given_info}
   - 구해야 할 것: {goal}

2. 단계별 계산 과정:
{calculation_steps}

3. 결론:
   위 계산 과정을 통해 정확한 답은 {result}입니다.
"""

# 최종 답변 생성을 위한 프롬프트 추가
final_answer_prompt = """
지금까지의 체계적인 추론 과정을 종합하여, 원래 질문에 답변하겠습니다:

{question}

추론 과정을 검토한 결과, 다음과 같은 결론에 도달했습니다:
{reasoning_conclusion}

따라서 최종 답변은:
{ANSWER_MARKER}
"""

# 수식 표시 문제 해결을 위한 설정
latex_delimiters = [
    {"left": "$$", "right": "$$", "display": True},
    {"left": "$", "right": "$", "display": False},
]


def reformat_math(text):
    """Gradio 구문(Katex)을 사용하도록 MathJax 구분 기호 수정.
    이것은 Gradio에서 수학 공식을 표시하기 위한 임시 해결책입니다. 현재로서는
    다른 latex_delimiters를 사용하여 예상대로 작동하게 하는 방법을 찾지 못했습니다...
    """
    text = re.sub(r"\\\[\s*(.*?)\s*\\\]", r"$$\1$$", text, flags=re.DOTALL)
    text = re.sub(r"\\\(\s*(.*?)\s*\\\)", r"$\1$", text, flags=re.DOTALL)
    return text


def user_input(message, history_original, history_thinking):
    """사용자 입력을 히스토리에 추가하고 입력 텍스트 상자 비우기"""
    return "", history_original + [
        gr.ChatMessage(role="user", content=message.replace(ANSWER_MARKER, ""))
    ], history_thinking + [
        gr.ChatMessage(role="user", content=message.replace(ANSWER_MARKER, ""))
    ]


def rebuild_messages(history: list):
    """중간 생각 과정 없이 모델이 사용할 히스토리에서 메시지 재구성"""
    messages = []
    for h in history:
        if isinstance(h, dict) and not h.get("metadata", {}).get("title", False):
            messages.append(h)
        elif (
            isinstance(h, gr.ChatMessage)
            and h.metadata.get("title", None) is None
            and isinstance(h.content, str)
        ):
            messages.append({"role": h.role, "content": h.content})
    return messages


def extract_info_from_question(question):
    """문제에서 주어진 정보와 목표를 추출합니다."""
    # 기본 값
    given_info = "문제에서 제공된 모든 조건과 수치"
    goal = "문제에서 요구하는 값이나 결과"
    
    # 일반적인 정보 추출 패턴
    if "몇 개" in question or "개수" in question:
        goal = "특정 조건을 만족하는 항목의 개수"
    elif "얼마" in question:
        goal = "특정 값 또는 금액"
    elif "나이" in question:
        goal = "사람의 나이"
    elif "확률" in question:
        goal = "특정 사건의 확률"
    elif "평균" in question:
        goal = "값들의 평균"
    
    return given_info, goal


@spaces.GPU
def bot_original(
    history: list,
    max_num_tokens: int,
    do_sample: bool,
    temperature: float,
):
    """원본 모델이 질문에 답변하도록 하기 (추론 과정 없이)"""

    # 나중에 스레드에서 토큰을 스트림으로 가져오기 위함
    streamer = transformers.TextIteratorStreamer(
        pipe.tokenizer,  # pyright: ignore
        skip_special_tokens=True,
        skip_prompt=True,
    )

    # 보조자 메시지 준비
    history.append(
        gr.ChatMessage(
            role="assistant",
            content=str(""),
        )
    )

    # 현재 채팅에 표시될 메시지
    messages = rebuild_messages(history[:-1])  # 마지막 빈 메시지 제외
    
    # 원본 모델은 추론 없이 바로 답변
    t = threading.Thread(
        target=pipe,
        args=(messages,),
        kwargs=dict(
            max_new_tokens=max_num_tokens,
            streamer=streamer,
            do_sample=do_sample,
            temperature=temperature,
        ),
    )
    t.start()

    for token in streamer:
        history[-1].content += token
        history[-1].content = reformat_math(history[-1].content)
        yield history
    t.join()

    yield history


@spaces.GPU
def bot_thinking(
    history: list,
    max_num_tokens: int,
    final_num_tokens: int,
    do_sample: bool,
    temperature: float,
):
    """추론 과정을 포함하여 모델이 질문에 답변하도록 하기"""

    # 나중에 스레드에서 토큰을 스트림으로 가져오기 위함
    streamer = transformers.TextIteratorStreamer(
        pipe.tokenizer,  # pyright: ignore
        skip_special_tokens=True,
        skip_prompt=True,
    )

    # 필요한 경우 추론에 질문을 다시 삽입하기 위함
    question = history[-1]["content"]
    
    # 문제에서 주어진 정보와 목표 추출
    given_info, goal = extract_info_from_question(question)

    # 보조자 메시지 준비
    history.append(
        gr.ChatMessage(
            role="assistant",
            content=str(""),
            metadata={"title": "🧠 생각 중...", "status": "pending"},
        )
    )

    # 현재 채팅에 표시될 추론 과정
    messages = rebuild_messages(history)
    
    # 전체 추론 과정을 저장할 변수
    full_reasoning = ""
    
    # 추론 과정에서 수집된 계산 단계 저장
    calculation_steps = ""
    current_conclusion = "아직 최종 결론에 도달하지 않았습니다."
    
    # 추론 단계 실행
    for i, prepend in enumerate(rethink_prepends):
        if i > 0:
            messages[-1]["content"] += "\n\n"
        
        # 첫 단계에서 일반적인 추론 가이드 추가
        if i == 0:
            messages[-1]["content"] += general_reasoning_guide + "\n\n"
        
        # 중간 단계에서 구조화된 추론 요약 추가
        if i > 1 and calculation_steps:
            structured_summary = structured_reasoning_prompt.format(
                given_info=given_info,
                goal=goal,
                calculation_steps=calculation_steps,
                current_conclusion=current_conclusion
            )
            messages[-1]["content"] += structured_summary + "\n\n"
        
        messages[-1]["content"] += prepend.format(question=question)

        t = threading.Thread(
            target=pipe,
            args=(messages,),
            kwargs=dict(
                max_new_tokens=max_num_tokens,
                streamer=streamer,
                do_sample=do_sample,
                temperature=temperature,
            ),
        )
        t.start()

        # 새 내용으로 히스토리 재구성
        if i == 0:
            history[-1].content += general_reasoning_guide + "\n\n"
        if i > 1 and calculation_steps:
            history[-1].content += structured_summary + "\n\n"
        history[-1].content += prepend.format(question=question)
        
        for token in streamer:
            history[-1].content += token
            history[-1].content = reformat_math(history[-1].content)
            yield history
        t.join()
        
        # 각 추론 단계의 결과를 full_reasoning에 저장
        full_reasoning = history[-1].content
        
        # 계산 단계 추출 및 업데이트
        new_content = history[-1].content.split(prepend.format(question=question))[-1]
        if "=" in new_content or ":" in new_content:
            # 계산 단계가 있는 것으로 간주
            calculation_steps += f"\n   - {new_content.strip()}"
        
        # 단계에서 가능한 결론 추출
        results = extract_calculation_results(new_content)
        if results:
            current_conclusion = f"현재 계산된 값: {results[-1]}"

    # 추론 완료, 이제 최종 답변을 생성
    history[-1].metadata = {"title": "💭 사고 과정", "status": "done"}
    
    # 추론 과정에서 도출된 모든 결과 추출
    all_results = extract_calculation_results(full_reasoning)
    
    # 결과가 있는 경우 검증 단계 추가
    if all_results and len(set(all_results)) > 1:
        # 결과별 빈도 계산
        result_counter = Counter(all_results)
        different_results = "\n".join([f"{result} (빈도: {freq}회)" for result, freq in result_counter.most_common()])
        
        # 최적의 결과 결정
        best_result = determine_best_result(all_results, full_reasoning)
        
        # 모델에게 가장 정확한 결과 선택 요청
        verify_prompt = verification_prompt.format(
            different_results=different_results,
            given_info=given_info,
            goal=goal,
            calculation_steps=calculation_steps,
            result=best_result
        )
        messages[-1]["content"] += "\n\n" + verify_prompt
        
        # 검증 단계 실행
        t = threading.Thread(
            target=pipe,
            args=(messages,),
            kwargs=dict(
                max_new_tokens=max_num_tokens // 2,
                streamer=streamer,
                do_sample=False,  # 확정적인 결과를 위해 샘플링 비활성화
                temperature=0.3,   # 낮은 온도 사용
            ),
        )
        t.start()
        
        history[-1].content += "\n\n" + verify_prompt
        for token in streamer:
            history[-1].content += token
            history[-1].content = reformat_math(history[-1].content)
            yield history
        t.join()
        
        # 검증 단계 후 full_reasoning 업데이트
        full_reasoning = history[-1].content
    
    # 최종 결과 결정
    final_results = extract_calculation_results(full_reasoning)
    best_result = determine_best_result(final_results, full_reasoning) if final_results else None
    
    # 최종 결론 생성
    if best_result is not None:
        reasoning_conclusion = f"추론 과정을 종합한 결과, 정확한 답변은 {best_result}입니다."
    else:
        # 결과를 추출할 수 없는 경우의 대비책
        reasoning_parts = full_reasoning.split("\n\n")
        reasoning_conclusion = "\n\n".join(reasoning_parts[-2:]) if len(reasoning_parts) > 2 else full_reasoning
    
    # 최종 답변 메시지 추가
    history.append(gr.ChatMessage(role="assistant", content=""))
    
    # 최종 답변을 위한 메시지 구성
    final_messages = rebuild_messages(history[:-1])  # 마지막 빈 메시지 제외
    final_prompt = final_answer_prompt.format(
        question=question,
        reasoning_conclusion=reasoning_conclusion,
        ANSWER_MARKER=ANSWER_MARKER
    )
    final_messages[-1]["content"] += "\n\n" + final_prompt
    
    # 최종 답변 생성
    t = threading.Thread(
        target=pipe,
        args=(final_messages,),
        kwargs=dict(
            max_new_tokens=final_num_tokens,
            streamer=streamer,
            do_sample=do_sample,
            temperature=temperature * 0.8,  # 최종 답변에 더 확신을 주기 위해 온도 약간 낮춤
        ),
    )
    t.start()
    
    # 최종 답변 스트리밍
    for token in streamer:
        history[-1].content += token
        history[-1].content = reformat_math(history[-1].content)
        yield history
    t.join()

    yield history


with gr.Blocks(fill_height=True, title="Vidraft ThinkFlow") as demo:
    # 제목과 설명
    gr.Markdown("# Vidraft ThinkFlow")
    gr.Markdown("### 추론 기능이 없는 LLM 모델의 수정 없이도 추론 기능을 자동으로 적용하는 LLM 추론 생성 플랫폼")
    
    with gr.Row(scale=1):
        with gr.Column(scale=2):
            gr.Markdown("## Before (Original)")
            chatbot_original = gr.Chatbot(
                scale=1,
                type="messages",
                latex_delimiters=latex_delimiters,
                label="Original Model (No Reasoning)"
            )
        
        with gr.Column(scale=2):
            gr.Markdown("## After (Thinking)")
            chatbot_thinking = gr.Chatbot(
                scale=1,
                type="messages",
                latex_delimiters=latex_delimiters,
                label="Model with Reasoning"
            )
    
    with gr.Row():
        # msg 텍스트박스를 먼저 정의
        msg = gr.Textbox(
            submit_btn=True,
            label="",
            show_label=False,
            placeholder="여기에 질문을 입력하세요.",
            autofocus=True,
        )
    
    # 예제 섹션 - msg 변수 정의 이후에 배치
    with gr.Accordion("EXAMPLES", open=False):
        examples = gr.Examples(
            examples=[
                "[출처: MATH-500)] 처음 100개의 양의 정수 중에서 3, 4, 5로 나누어 떨어지는 수는 몇 개입니까?",
                "[출처: MATH-500)] 잉크의 땅에서 돈 시스템은 독특합니다. 트링킛 1개는 블링킷 4개와 같고, 블링킷 3개는 드링크 7개와 같습니다. 트링킷에서 드링크 56개의 가치는 얼마입니까?",
                "[출처: MATH-500)] 에이미, 벤, 크리스의 평균 나이는 6살입니다. 4년 전 크리스는 지금 에이미와 같은 나이였습니다. 4년 후 벤의 나이는 그때 에이미의 나이의 $\\frac{3}{5}$가 될 것입니다. 크리스는 지금 몇 살입니까?",
                "[출처: MATH-500)] 노란색과 파란색 구슬이 들어 있는 가방이 있습니다. 현재 파란색 구슬과 노란색 구슬의 비율은 4:3입니다. 파란색 구슬 5개를 더하고 노란색 구슬 3개를 제거하면 비율은 7:3이 됩니다. 더 넣기 전에 가방에 파란색 구슬이 몇 개 있었습니까?",
                "수학 동아리에서 다가올 여행을 위한 기금 모금을 위해 베이킹 세일을 열고 있습니다. 3개에 54달러짜리 쿠키를 1달러에 판매하고, 20개에 컵케이크를 각각 2달러에 판매하고, 35개에 브라우니를 각각 1달러에 판매합니다. 수학 동아리에서 이 제품을 굽는 데 15달러가 들었다면, 수익은 얼마였을까요?"                
            ],
            inputs=msg
        )
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("""## 매개변수 조정""")
            num_tokens = gr.Slider(
                50,
                4000,
                2000,
                step=1,
                label="추론 단계당 최대 토큰 수",
                interactive=True,
            )
            final_num_tokens = gr.Slider(
                50,
                4000,
                2000,
                step=1,
                label="최종 답변의 최대 토큰 수",
                interactive=True,
            )
            do_sample = gr.Checkbox(True, label="샘플링 사용")
            temperature = gr.Slider(0.1, 1.0, 0.7, step=0.1, label="온도")

    # 사용자가 메시지를 제출하면 두 봇이 동시에 응답합니다
    msg.submit(
        user_input,
        [msg, chatbot_original, chatbot_thinking],  # 입력
        [msg, chatbot_original, chatbot_thinking],  # 출력
    ).then(
        bot_original,
        [
            chatbot_original,
            num_tokens,
            do_sample,
            temperature,
        ],
        chatbot_original,  # 출력에서 새 히스토리 저장
    ).then(
        bot_thinking,
        [
            chatbot_thinking,
            num_tokens,
            final_num_tokens,
            do_sample,
            temperature,
        ],
        chatbot_thinking,  # 출력에서 새 히스토리 저장
    )

if __name__ == "__main__":
    demo.queue().launch()  # title 매개변수 제거