from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Union, Dict, Any
from router.bale.bale_buttons import BUTTON_TEXT_TO_CALLBACK_LIST
import json
from typing import Optional, Callable, List, Any
from pydantic import BaseModel


class SingleSearchData(BaseModel):
    score:float
    id:str
    content:str

class BMNewSemanticSearchOutput(BaseModel):
    query: str
    result : List[SingleSearchData]
    metadata : Dict
    embed_model_time : float
    cosine_similarity_time : float
    rerank_time : float
    
    
class DbRule(BaseModel):
    rule_id: str
    rule_content: str
    rule_type: str
    section_id: str
    section_content: str
    section_full_path :str
    qanon_id: str
    qanon_etebar: str
    qanon_title: str
    state_etebar: str

class InputRule(BaseModel):
    rule_id: str
    rule_content: str
    rule_type: str
    section_id: str
    section_content: str

class SemanticSearchP2P(BaseModel):
    in_rule: InputRule
    db_rule: DbRule
    score: float = 0
    metadata: Dict

class BaleStartMessageForm(BaseModel):
    id: int
    is_bot: bool = False
    first_name: str
    last_name: Optional[str] = None
    username: Optional[str] = None


class BaleStartMessageChat(BaseModel):
    id: int
    type: str
    username: Optional[str] = None
    first_name: Optional[str] = None


class BaleStartMessage(BaseModel):
    message_id: int
    from_user: BaleStartMessageForm = Field(..., alias="from")
    date: int
    chat: BaleStartMessageChat
    text: str
    entities: List[dict] = []

    class Config:
        populate_by_name = True


class BaleCallbackFrom(BaseModel):
    id: int
    is_bot: bool
    first_name: str
    username: Optional[str] = None


class BaleCallbackMessage(BaseModel):
    message_id: int
    chat: BaleStartMessageChat
    text: Optional[str]


class BaleCallbackQuery(BaseModel):
    id: str
    from_user: BaleCallbackFrom = Field(..., alias="from")
    message: BaleCallbackMessage
    data: str

    class Config:
        populate_by_name = True


class BaleUpdate(BaseModel):
    update_id: int
    message: Optional[BaleStartMessage] = None
    callback_query: Optional[BaleCallbackQuery] = None

    class Config:
        exclude_none = True


class QaChatSingle(BaseModel):
    id: str
    chat_id: int
    user_query: str
    model_key: str
    model_effort: str

    retrived_passage: str
    retrived_ref_ids: str
    retrived_duration: Optional[int] = 0
    prompt_type: str = "question"
    llm_duration: int
    full_duration: Optional[int] = 0
    time_create: Optional[int] = 0
    used_ref_ids: Optional[str] = ""
    status_text: Optional[str] = ""
    status: Optional[int] = 0
    prompt_answer: str
    other_info: dict | None


class QaChatBlock(BaseModel):
    id: str
    title: str
    user_id: str
    is_premium: bool
    chat: QaChatSingle
    total_token: int
    is_end: bool


class QaChat(BaseModel):
    time_stamp: int = 0
    id: str
    chat_id: int
    title: Optional[str] = ""
    user_id: str
    user_query: str
    query_type: str = (
        "question"  # llm -> greeting, other, legal_question | rag -> question
    )
    full_duration: Optional[float] = 0
    other_info: Optional[dict] = ""

    ss_ref_ids: Optional[List[str]] = ""
    ss_model_key: Optional[str] = ""
    ss_duration: Optional[float] = 0
    ss_answer: Optional[str] = ""

    llm_ref_ids: Optional[List[str]] = []
    llm_model_key: Optional[str] = ""
    llm_duration: Optional[float] = 0
    llm_answer: Optional[str] = ""

    status_text: Optional[str] = ""
    status: Optional[int] = 0


class ConflictDetection(BaseModel):
    explanation_of_conflict: str
    has_confict: bool = False


class ConflictTypeDetection(BaseModel):
    conflict_type: Literal["Doctrinal_conflict", "Practical_conflict"]
    explanation_of_type: str


class RelationIdentification(BaseModel):
    reasoning: str
    relation_type: Literal[
        "مطلق مقدم، مقید موخر",
        "مقید مقدم، مطلق موخر",
        "تکرار حکم",
        "تعارض مستقر",
        "بدون تعارض",
    ]
class Evaluation(BaseModel):
    is_subject_unity_assessment_correct: bool
    is_conflict_detection_correct: bool
    is_conflict_type_detection_correct: bool
    is_relation_type_detection_correct: bool
    valid_relation_type: str
    comments: str

class SingleRuleRelation(BaseModel):
    rule_id: str
    rule_content: str
    rule_type: str

    section_id: str
    section_content: str

    qanon_etebar: Optional[str] = None
    qanon_id: Optional[str] = None
    qanon_title: Optional[str] = None

    state_etebar: Optional[str] = None

    date: Optional[str] = None

    def __str__(self):
        dict_ = {
            # "rule_id": self.rule_id,
            "متن حکم": self.rule_content,
            "نوع حکم": self.rule_type,
            "متن ماده قانونی": self.section_content,
        }
        # Return the dictionary as a formatted JSON string
        return json.dumps(dict_, ensure_ascii=False, indent=2)


class SubjectUnity(BaseModel):
    has_subject_unity: Literal["yes", "no", "yes_under_assumptions"]
    required_assumptions: Optional[str] = None
    reasoning: Optional[str]


class RuleRelation(BaseModel):
    # ----- input rule
    in_rule: SingleRuleRelation

    # ----- semantic-search close rule
    db_rule: SingleRuleRelation

    # ----- subject-unity data
    subject_unity: SubjectUnity | None = None

    # ----- conflict-detection data
    conflict_detection: ConflictDetection | None = None

    # ----- conflict-type-detection data
    conflict_type_detection: ConflictTypeDetection | None = None

    # ----- relation-identification data
    relation_identification: RelationIdentification | None = None

class StateDetail(BaseModel):
    state: str 
    message: str
    button_text: str
    end_buttons: List = []
    inline_buttons: Optional[List[List[dict]]] = None
    
    handler : str = None
    allow_empty_input: bool = True  

class BaleUser(BaseModel):
    uc_id: str
    chat_id: int
    user_id: int
    update: BaleUpdate
    username: str = None
    is_bot: bool = False
    is_vip: bool = False
    first_name: str = ""
    last_name: str = ""
    message_limit:int = 0
    rule_relation: RuleRelation | None = None
    subject_unities:Dict = {}

    # ---- defaults
    effort: Literal["medium", "low"] = "low"
    limit: int = 10
    permission: str = "normal"

    # ---- runtime
    is_processing_lock : bool = False
    is_call_back_query : bool = False
    state_detail : StateDetail = None
    active_message_id : int = 0
    
    input_query: str = ""  # ورودی  کاربر
    last_input_query: str = ""  # ورودی  کاربر
    call_back_query: str = ""  # ورودی  کاربر
    _query_type: str = ""  # ورودی  کاربر
    sub_state: str = ""  # برای روندی ها

    all_qq: bool = False

    # ---- memory
    last_query: Dict[str, str] = {}  # آخرین سوال کاربر
    last_result: Dict[str, Any] = {}  # آخرین نتایج
    last_runtime: Dict[str, Any] = (
        {}
    )  # مثلا {"GP": {"effort": "medium", "limit": 20}} برای بقیه چیزها
    stack: Dict[str, List] = {}  # پشته برای داده های مرحله ای


class KeyboardItem(BaseModel):
    text: str


class InlineKeyboardItem(BaseModel):
    text: str
    callback_data: str


class ReplyMarkup(BaseModel):
    keyboard: List[List[Union[KeyboardItem, Dict]]] = Field(
        default_factory=lambda: BUTTON_TEXT_TO_CALLBACK_LIST
    )
    resize_keyboard: bool = True
    one_time_keyboard: bool = False
    inline_keyboard: List[List[Union[InlineKeyboardItem, Dict]]] | None = (
        None  # دکمه ها ی داخال صفحه
    )


class BalePayload(BaseModel):
    reply_markup: ReplyMarkup = ReplyMarkup()
    chat_id: int
    message_id: int = None
    text: str


class Step(BaseModel):
    name: str
    level : int
    data_input : str
    data_output : str
    bot_output : List