# Helper functions
import re
import json
from typing import List, Dict, Union

from models.memory import SimpleConversationMemory

def format_chat_history(memory_obj: SimpleConversationMemory):
    """Format chat history for display using custom memory"""
    messages = memory_obj.get_messages()
    if not messages:
        return "Không có lịch sử hội thoại trước."
    
    formatted = []
    for msg in messages:
        role = "User" if msg["role"] == "user" else "Assistant"
        formatted.append(f"{role}: {msg['content']}")
    
    return "\n".join(formatted)

def preprocess_related_questions(related_questions_input: Union[str, List[Dict[str, str]]]) -> List[Dict[str, str]]:
    fallback_questions = [
        {"question": "Quy định pháp luật Việt Nam hiện hành về xử lý tranh chấp hợp đồng dân sự được quy định trong văn bản nào?"},
        {"question": "Trường hợp nào thì một bản án có thể được sử dụng làm án lệ theo quy định của pháp luật Việt Nam?"},
        {"question": "Các nguyên tắc cơ bản của Bộ luật Dân sự Việt Nam năm 2015 được quy định tại điều khoản nào?"},
        {"question": "Nghị định nào quy định về xử phạt vi phạm hành chính trong lĩnh vực hôn nhân và gia đình tại Việt Nam?"},
    ]
    
    if isinstance(related_questions_input, str):
        cleaned_input = re.sub(r'^```json\s*|\s*```$', '', related_questions_input).strip()
        try:
            related_questions = json.loads(cleaned_input)
        except json.JSONDecodeError:
            return fallback_questions[:5]
    else:
        related_questions = related_questions_input
    
    if not isinstance(related_questions, list):
        return fallback_questions[:5]
    
    valid_questions = [
        q for q in related_questions
        if isinstance(q, dict) and "question" in q and isinstance(q["question"], str) and q["question"].strip()
    ]
    
    seen = set()
    unique_questions = []
    for q in valid_questions:
        question_text = q["question"].strip()
        if question_text not in seen:
            seen.add(question_text)
            unique_questions.append({"question": question_text})
    
    legal_keywords = r"(Luật|Bộ luật|Nghị định|Thông tư|Quy định|án lệ|Việt Nam|tòa án|pháp luật|điều luật|Bảo hiểm xã hội)"
    filtered_questions = [
        q for q in unique_questions
        if re.search(legal_keywords, q["question"], re.IGNORECASE)
    ]
    
    if len(filtered_questions) < 4:
        for fq in fallback_questions:
            if len(filtered_questions) >= 4:
                break
            if fq["question"] not in seen:
                filtered_questions.append(fq)
                seen.add(fq["question"])
    
    return filtered_questions[:4]