from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
import time, os, traceback
from .base_model import Query, LLMOutput, LLMInput, Title
from .ai_data_parser import AsyncCore
from .chatbot_handler import (
    InitHybridRetrieverReranker,
    format_answer_bale,
    get_user_prompt2,
    get_user_prompt3,
    load_faiss_index,
    get_in_form,chunked_simple_text
)
from .static import (
    EMBED_MODEL_PATH,
    FAISS_INDEX_PATH,
    FAISS_METADATA_PATH,
    LLM_URL,
    SYSTEM_PROMPT_FINALL,
    RERANKER_MODEL_PATH,
    LLM_ERROR,
    MODEL_KEY,
    MODEL_NAME,
    OUTPUT_PATH_LLM,
    REASONING_EFFORT,
    TASK_NAME,
    LLM_TIME_OUT,
    MAX_TOKEN,
    SYSTEM_PROPMT2,
)


# ################################################## Global-params
router = APIRouter(tags=["ragchat"])
# # settings= get_settings()


METADATA_DICT, FAISS_INDEX = load_faiss_index(
    index_path=FAISS_INDEX_PATH, metadata_path=FAISS_METADATA_PATH
)

RAG = InitHybridRetrieverReranker(
    embeder_path=EMBED_MODEL_PATH,
    reranker_path=RERANKER_MODEL_PATH,
    dict_content=METADATA_DICT,
    faiss_index=FAISS_INDEX,
    dense_alpha=0.6,
    device="cuda",
)

RUNNER_PROMPT = AsyncCore(
    model_name=MODEL_NAME,
    api_url=LLM_URL,
    output_path=OUTPUT_PATH_LLM,
    task_name=TASK_NAME,
    output_schema=LLMOutput,
    reasoning_effort=REASONING_EFFORT,
    ai_code_version=MODEL_KEY,
    request_timeout=LLM_TIME_OUT,
    max_token=MAX_TOKEN,
    save_number=1,
)

from pydantic import BaseModel

SYS_PROMPT = """
شما یک ربات متخصص حقوقی و قانونی هستید به نام «قانون‌یار». 
هویت شما: یک وکیل باتجربه، دقیق، مودب و مشتاق کمک‌کردن به افراد در فهم مسائل حقوقی.

در این مرحله، وظیفه شما «تشخیص نوع پیام» است؛ اما باید یک دلیل کوتاه (answer) بدهید
که لحنش لحن یک وکیل کارکشته و حامی باشد.

برچسب خروجی فقط یکی از این موارد است:
- legal_question
- greeting
- other

تعریف‌ها:
1. legal_question: هر پرسش یا متن مرتبط با قوانین، مقررات، دادگاه، وکالت، قرارداد، مالیات، ارث، نکاح و طلاق، شکایت، چک، سفته، ثبت شرکت، حقوق کار و هر موضوع مشابه.
2. greeting: پیام‌های غیرحقوقی شامل سلام، معرفی، احوال‌پرسی، تشکر، خداحافظی.
3. other: پیام‌هایی که خارج از دو دسته بالا قرار می‌گیرند.

نکات مهم:
- answer باید بسیار کوتاه باشد (حداکثر دو خط).
- اگر پیام حتی کمی رنگ‌وبوی حقوقی داشته باشد، answer را به سمت توضیح حقوقی ببرید.
- در حالت greeting با لحن رسمی و دوستانه یک وکیل پاسخ بدهید.
- در حالت other، در عین محترمانه‌گفتن اینکه پیام حقوقی نیست، سعی کنید یک زاویهٔ حقوقی احتمالی مطرح کنید تا راه کمک باز بماند.

خروجی شما حتماً باید شامل دو فیلد باشد:
label: یکی از سه برچسب
answer: توضیح کوتاه، دوستانه، مودب و با لحن یک وکیل حامی

خروجی فقط JSON باشد.
"""


ASSIST_PROMPT = """
### Example 1:
input:
"عین کلی"
output:
{
    "label": "legal_question",
    "answer": "موضوع شما به یکی از مفاهیم مهم حقوق مدنی مربوط می‌شود و با کمال میل راهنمایی‌تان می‌کنم."
}

### Example 2:
input:
"طرز تهیه ماست"
output:
{
    "label":"other",
    "answer":"موضوع شما ماهیت حقوقی ندارد، اما اگر دغدغه‌تان مرتبط با استاندارد یا مسئولیت تولید باشد میتوانم به شما کمک کنم، می شود سوال خود را بهتر و کامل تر بپرسید؟ میخواهم دقیق بدانم چه قسمتی از موارد حقوقی تولید ماست رو میخواهید بدونید"
}

### Example 3:
input:
"سلام، تو کی هستی؟"
output:
{
    "label":"greeting",
    "answer":"سلام و ارادت؛ من قانون‌یار هستم، وکیل همراه شما برای هر پرسش حقوقی."
}
"""

class ChatOutput(BaseModel):
    label: str
    answer: str | None = None
    
RUNNER_CHAT = AsyncCore(
    model_name=MODEL_NAME,
    api_url=LLM_URL,
    output_path="/home2/rag_qa_chat2/data/_temp/",
    task_name="type-of-chat",
    output_schema=ChatOutput,
    reasoning_effort="low",
    ai_code_version=MODEL_KEY,
    request_timeout=LLM_TIME_OUT,
    max_token=256,
    save_number=1,
    top_p=0.8,
    temperature=0.5
)

# legal_question
# greeting
# other

# label
# reasoning
async def chat_bot_run(query):
    try:
        s0 = time.time()
        print(f'query {query}')
        type_, _ = await RUNNER_CHAT.single_simple_async_proccess_item(
            item={
                "user_prompt": query,
                "system_prompt": SYS_PROMPT,
                "assistant_prompt":ASSIST_PROMPT
            }
        )
        type_chat =  type_.get('label', None)
        answer =  type_['answer']
        s = time.time()
        print(f'type_llm {type_} {s-s0}\'s Toke')
        if type_chat != "legal_question":
            return chunked_simple_text(answer) 
        
        sections_dict = await RAG.search_base(
            query,
            final_k=10,
            topk_dense=100,
            topk_sparse=100,
            pre_rerank_k=100,
        )
        e = time.time()
        # input_data = LLMInput(query=query, knowledge=sections_dict)
        # prompt = get_user_prompt2(input_data)
        prompt = get_user_prompt3(query=query, knowledge_json=sections_dict)

        llm_answer, _ = await RUNNER_PROMPT.single_simple_async_proccess_item(
            item={"user_prompt": prompt, "system_prompt": SYSTEM_PROPMT2},
        )
        ee = time.time()
        finall = format_answer_bale(
            answer_text=llm_answer["text"], sources=llm_answer["source"]
        )
        eee = time.time()
        print(f"Rag = {e-s}", f"llm_answer = {ee-e}", f"Format = {eee-ee}", sep="\n")
        return finall
    except:
        traceback.print_exc()


class RagInput(BaseModel):
    query: str
    limit:int = 10
    
class RagOutput (BaseModel):
    finall_form: str
    duration: int
    ids: List[str]

async def rag_run(input:RagInput):
    try:
        
        s = time.time()
        sections_dict = await RAG.search_base(
            query=input.query,
            final_k=input.limit,
            topk_dense=100,
            topk_sparse=100,
            pre_rerank_k=100,
        )
        e = time.time()
        finall = get_in_form(title=input.query, sections=sections_dict)
        ee = time.time()
        print(f"Rag = {e-s}", f"Form = {ee-e}", sep="\n")
        finall = RagOutput(

        )
        return finall
    except:
        traceback.print_exc()


@router.post("/run_chat")
async def run_chat(payload: Query, request: Request):
    s = time.time()
    try:
        answer = await chat_bot_run(payload.query)
    except:
        print(f"chat_bot_run FAIL!")
        answer = LLM_ERROR
    e = time.time()
    print(f"Total Time {e-s:.2f}'s")
    return JSONResponse({"result": answer}, status_code=201)


@router.post("/run_rag")
async def run_chat(payload: Query, request: Request):
    s = time.time()
    try:
        answer = await rag_run(payload.query)
    except:
        print(f"chat_bot_run FAIL!")
        answer = LLM_ERROR
    e = time.time()
    print(f"Total Time {e-s:.2f}'s")
    return JSONResponse({"result": answer}, status_code=201)