from data import debug_print,df,eligibility_df,llm1
from nodes.intent import get_pretty_state_string,CreditCardState
from pydantic_schema import CreditCardRecommendation
from langchain_core.messages import HumanMessage,AIMessage
import pprint

#formatting and structuring the final response
def extract_card_info_combined(agent_response: CreditCardRecommendation) -> str:
    try:
        best_card = agent_response.best_card
        explanation_list = agent_response.explanation
        explanation_list_points = [f"<li style='color: black;'>{point}</li>" for point in explanation_list]
        explanation = " ".join(explanation_list_points)
        explanation = "<ul> " + explanation + " </ul>"

        debug_print("UTIL", f"Extracted best card: '{best_card}'")
        debug_print("UTIL", f"Explanation length: {len(explanation)}")
    except Exception as e:
        debug_print("ERROR", f"Failed to parse Gemini response as JSON: {str(e)}")
        best_card = "N/A"
        explanation = "No explanation provided."

    best_card_block = (
        f"<strong style='color: black;'>Best Card:</strong> {best_card}\n\n"
        f"<strong style='color: black;'>Why It's The Best:</strong> \n{explanation}"
    )

    return best_card_block

def build_card_rows(card_names):
    debug_print("UTIL", f"Building card rows for {len(card_names)} cards")

    card_rows = []
    card_links = []
    for name in card_names:
        match = df[df["name"] == name]
        joining_fee = "N/A"
        annual_fee = "N/A"
        issuer_link = "N/A"
        description = ""

        if name in eligibility_df["Name"].values:
            row = eligibility_df[eligibility_df["Name"] == name].iloc[0]
            joining_fee = str(row.get("Joining fee", "N/A"))
            annual_fee = str(row.get("Annual fee", "N/A"))
            issuer_link = row.get("Issuer Link", "N/A")
        if not match.empty:
            description = match.iloc[0].get("description", "")
        card_rows.append([name, joining_fee, annual_fee, description])
        card_links.append(issuer_link)

    debug_print("UTIL", f"Built {len(card_rows)} card rows")
    return card_rows, card_links


async def format_output_node(state: CreditCardState):
    debug_print("NODE", f"Entering format_output_node with state:\n {get_pretty_state_string(state)}\n")
    top_card_html = '''
        <div style="
            background-color: #fff3e0;
            color: #212121;
            border-radius: 16px;
            padding: 20px;
            border: 2px solid #ffa726;
            box-shadow: 2px 2px 8px rgba(0,0,0,0.1);
            margin-bottom: 16px;
            font-family: sans-serif;
            font-size: 8px;
        ">
            <pre style="white-space: pre-wrap; font-size: 13px; color: #212121;">{message}</pre>
        </div>
    '''
    if not state.get("ranked_cards"):
        debug_print("NODE", "No eligible cards in ranked_cards. Skipping agent interpretation.")
        message = "There are no eligible cards available based on your profile at this time."
        
        return {
            "top_card_html": "",  
            "top_card": "",     
            "top_card_description": [message],  
            "card_rows": [],
            "card_names": [],
            "card_lookup": {},
            "card_links": []
        }

    
    if "messages" not in state or not state["messages"]:
        debug_print("NODE", "No messages found in state, returning empty output")
        return {
            "top_card_html": top_card_html.format(message="No messages found"),
            "top_card": "No messages found",
            "top_card_description": [],
            "card_rows": [],
            "card_names": [],
            "card_lookup": {}
        }
            
    last_message = state["messages"][-1]

    if isinstance(last_message, AIMessage):
        print(f"Agent: {last_message.content}")
        
        model = llm1

        prompt_template = """
            **SYSTEM ROLE**
            You are an expert content transformation agent. Your task is to analyze the final message from an AI assistant and create a structured JSON response.
            
            **CONTEXT**
            - The user's original request was: "{user_query}"
            - The AI assistant's final message is:
            {message}
            
            **EXTRACTION RULES**
            You must follow these rules precisely:
            1.  **Analyze Outcome:** Determine if a specific card was recommended (SUCCESS_CASE) or if no suitable card was found (FAILURE_CASE).
            2.  **SUCCESS_CASE (A card was recommended):**
                - Set `card_found` to `True`.
                - Extract the exact name of the recommended card for `best_card`.
                - Structure the reasons and benefits as a list of strings for the `explanation` field.
                - `reply_if_card_not_found` MUST be `null`.
            3.  **FAILURE_CASE (No single card was recommended):**
                - Set `card_found` to `False`.
                - Rephrase the assistant's message into a single, user-friendly `reply_if_card_not_found`.
                - `best_card` and `explanation` fields MUST be `null`.
            
            **EXAMPLE**
            - IF the AI_Agent_Message is: "The user's goal is to minimize debt... I can suggest exploring options like the Axis Bank Burgundy Private Credit Card... Another option could be the SBI SimplySAVE Credit Card..."
            - THEN the `reply_if_card_not_found` field in your JSON should be: "While I couldn't find a single credit card that perfectly aligns with your goal of minimizing debt, my research identified a couple of different approaches you could consider. For those with a high net worth, premium cards like the Axis Bank Burgundy Private might offer very low interest rates... For a more accessible option, cards like the SBI SimplySAVE... I suggest exploring these two types of cards to see which strategy best fits your financial profile."
            
            **FINAL INSTRUCTION**
            Your entire response MUST be a valid JSON object that conforms to the required schema. Do not add any other text or formatting.
            """

        # Extracting the JSON schema from your Pydantic model
        json_schema = CreditCardRecommendation.model_json_schema()
        
        prompt = prompt_template.format(user_query=state.get("raw_query", ""), message=last_message.content)
        
        try:
            response = await model.ainvoke(
                [HumanMessage(content=prompt)],
                extra_body={"guided_json": json_schema}
            )
            debug_print("STRUCTURED_OUTPUT", f"Raw JSON string from LLM: {response.content}")
            
            # Parsing the JSON string response and creating the Pydantic object
            structured_response = CreditCardRecommendation.model_validate_json(response.content)
            debug_print("STRUCTURED_OUTPUT", "Successfully parsed into Pydantic object:")
            pprint.pprint(structured_response.model_dump(), indent=2)

        except Exception as e:
            debug_print("ERROR", f"Error invoking LLM with structured output: {str(e)}")
            return {
                "top_card_html": top_card_html.format(message="Error processing AI response"),
                "top_card": "Error processing AI response",
                "top_card_description": [],
                "card_rows": [],
                "card_names": [],
                "card_lookup": {}
            }
        
        if not structured_response.card_found:
            user_reply = structured_response.reply_if_card_not_found or "Unfortunately, a suitable card could not be found for your specific query"
            debug_print("NODE", f"No specific card found. Generated User Reply: {user_reply}")
            final_html_output = top_card_html.format(message=user_reply)
            return {
                "top_card_html": final_html_output,
                "top_card": user_reply,
                "top_card_description": [],
                "card_rows": [],
                "card_names": [],
                "card_lookup": {}
            }
        
        extracted_response = extract_card_info_combined(structured_response)
        card_names = [card["name"] if isinstance(card, dict) else card for card in state["ranked_cards"]]
        debug_print("NODE", f"Using {len(card_names)} card names from ranked_cards")
        top_card_html = f"""
            <div style="
                background-color: #fff3e0;
                color: #212121;
                border-radius: 16px;
                padding: 20px;
                border: 2px solid #ffa726;
                box-shadow: 2px 2px 8px rgba(0,0,0,0.1);
                margin-bottom: 16px;
                font-family: sans-serif;
                font-size: 8px;
            ">
                <pre style="white-space: pre-wrap; font-size: 13px; color: #212121;">{extracted_response}</pre>
            </div>
        """
        card_rows, card_links = build_card_rows(card_names)
        card_lookup = {row[0]: row[3] for row in card_rows}  
        debug_print("NODE", f"Built {len(card_rows)} card rows")

        return {
            "top_card_html": top_card_html.format(message=extracted_response),
            "top_card": structured_response.best_card,
            "top_card_description": structured_response.explanation,
            "card_rows": card_rows,
            "card_names": card_names,
            "card_lookup": card_lookup,
            "card_links": card_links
        }

    debug_print("NODE", "Last message was not from AI, skipping formatting.")
    return {}