"""
Interactive Advisor Bot — The pedagogical chat loop
"""

import json
import urllib.request
import urllib.parse
from datetime import datetime
from langchain_groq import ChatGroq
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from langgraph.prebuilt import create_react_agent
from langchain_core.tools import tool
from typing import Optional

from config import get_settings
from services.database import get_random_scenario

# Store active sessions in memory
_sessions: dict[str, list] = {}

@tool
def SerpApi_Search(query: str) -> str:
    """Searches the web for current prices, news, and market data."""
    settings = get_settings()
    if not settings.SERPAPI_API_KEY:
        return "[SerpApi Error] API key not configured."
    
    url = f"https://serpapi.com/search.json?q={urllib.parse.quote(query)}&api_key={settings.SERPAPI_API_KEY}"
    try:
        req = urllib.request.Request(url)
        with urllib.request.urlopen(req) as response:
            data = json.loads(response.read().decode())
            
            result = []
            if "answer_box" in data:
                ans = data['answer_box'].get('answer') or data['answer_box'].get('snippet') or ""
                if ans:
                    result.append(f"Answer Box: {ans}")
            
            if "organic_results" in data:
                for res in data["organic_results"][:3]: # Top 3 results
                    result.append(f"Title: {res.get('title')}\nSnippet: {res.get('snippet')}")
            
            if not result:
                return f"No results found for '{query}'"
                
            return "\n\n".join(result)
    except Exception as e:
        return f"[SerpApi Error] {str(e)}"

@tool
def mcq_scenarios() -> str:
    """Fetches a random investment scenario from the database."""
    scenario = get_random_scenario()
    if not scenario:
        return json.dumps({"error": "No scenarios found in the database. Please generate some first!"})
    
    return json.dumps({
        "title": scenario["title"],
        "scenario_paragraph": scenario["scenario_paragraph"],
        "risk_level": scenario["risk_level"],
        "best_answer": scenario["best_answer"],
        "best_answer_rationale": scenario["best_answer_rationale"],
        "other_options": [
            {
                "option": scenario["other_option1"],
                "rationale": scenario["other_option1_exp"]
            },
            {
                "option": scenario["other_option2"],
                "rationale": scenario["other_option2_exp"]
            },
            {
                "option": scenario["other_option3"],
                "rationale": scenario["other_option3_exp"]
            }
        ]
    })

def get_chat_response(session_id: str, user_message: str, knowledge_base_id: Optional[str] = None) -> dict:
    """Handle a user's chat message and manage the educational loop via Agent."""
    
    from services.rag_engine import search_kb

    @tool
    def search_knowledge_base(query: str) -> str:
        """Searches the user-selected Knowledge Base for financial documents, PDFs, and definitions."""
        if not knowledge_base_id:
            return "No Knowledge Base is currently attached. Do not attempt to search it again."
        return search_kb(knowledge_base_id, query)
    
    settings = get_settings()
    llm = ChatGroq(
        api_key=settings.GROQ_API_KEY,
        model_name=settings.GROQ_MODEL,
        temperature=0.7,
    )
    
    today = datetime.now().strftime("%Y-%m-%d")
    
    system_prompt_str = f"""Role: Expert Investment Advisor AI for markets, strategy, and portfolio education.

1. INVESTMENT ADVICE & REAL-TIME DATA

Scope: Stocks, bonds, ETFs, diversification, risk, and analysis.

Tool: Use SerpApi_Search for all news, current prices (e.g., Gold in Egypt), and market data.

No-Refusal Rule: You have LIVE data access via SerpAPI. If asked for current info, IMMEDIATELY use the tool. NEVER say you cannot access real-time data.

Date: Today is {today}.

2. MCQ GENERATION (Practice)

Format: Question + 4 options (A-D).

Single-Line Rule: Write the question and each option into one continuous line each including the answers should be in bullet points.

Logic: Wait for user answer. Then provide the correct answer, detailed feedback for all choices, and educational context.

**. SCENARIO PRESENTATION (UI/UX):**
When a scenario is retrieved, present it with high readability using this exact structure:
---
### 📊 Investment Case Study
> [Insert a concise, professional paragraph describing the situation.]

**Key Market Data:**
* 💵 **Initial Capital:** [Value]
* 📈 **Asset Class:** [Type]
* ⏱️ **Time Horizon:** [Duration]
* ⚠️ **Risk Level:** [Rating]

**Select the Best Course of Action:**
* **A)** [Option A text]
* **B)** [Option B text]
* **C)** [Option C text]
* **D)** [Option D text]
---
*Instruction: Wait for the user's letter (A-D) before providing the rationale.*

3. MCQ SCENARIO QUESTIONS (Database)

Tool: Use mcq_scenarios.

Flow: Present scenario/table → Show randomized A-D options → Wait for answer → Provide feedback using bestAnswer or otherAnswer rationales.

4. CONSTRAINTS

Investment Only: Strictly ignore non-investment topics (politics, private life, etc.).

Evidence-Based: Support all advice with data from SerpApi_Search.

Clarity: Explain complex financial terms in accessible language."""

    tools = [SerpApi_Search, mcq_scenarios, search_knowledge_base]
    agent_executor = create_react_agent(llm, tools, prompt=system_prompt_str)
    
    if session_id not in _sessions:
        _sessions[session_id] = []
        
    chat_history = _sessions[session_id]
    
    try:
        # Pass the accumulated history plus the new user message
        messages_to_pass = chat_history + [HumanMessage(content=user_message)]

        rag_used = False
        # Try with tools first; retry without tools if Groq fails to generate a valid tool call
        answer_str = ""
        for attempt in range(2):
            try:
                if attempt == 0:
                    response = agent_executor.invoke({
                        "messages": messages_to_pass
                    })
                    # Check if the RAG tool was called
                    for msg in response["messages"]:
                        if isinstance(msg, ToolMessage) and msg.name == "search_knowledge_base":
                            rag_used = True
                else:
                    # Retry: direct LLM call without tools (fallback)
                    print("[Chat Agent] Retrying without tools (fallback)...")
                    fallback_response = llm.invoke(messages_to_pass)
                    answer_str = fallback_response.content if isinstance(fallback_response.content, str) else str(fallback_response.content)
                    break

                # Robustly extract text from the agent's response
                for msg in reversed(response["messages"]):
                    if isinstance(msg, AIMessage):
                        content = msg.content
                        if isinstance(content, str) and content.strip():
                            answer_str = content.strip()
                            break
                        elif isinstance(content, list):
                            parts = []
                            for block in content:
                                if isinstance(block, dict):
                                    parts.append(block.get("text", ""))
                                elif isinstance(block, str):
                                    parts.append(block)
                            joined = "".join(parts).strip()
                            if joined:
                                answer_str = joined
                                break
                if answer_str:
                    break
            except Exception as inner_e:
                err_str = str(inner_e).lower()
                if "tool_use_failed" in err_str or "failed_generation" in err_str or ("400" in err_str and "function" in err_str):
                    print(f"[Chat Agent] Tool call failed (attempt {attempt+1}): {inner_e}")
                    if attempt == 0:
                        continue  # retry without tools
                raise  # re-raise other errors
        
        if not answer_str:
            answer_str = "I apologize, but I couldn't formulate a response. Please try asking in a different way."
            
        # update history
        chat_history.append(HumanMessage(content=user_message))
        chat_history.append(AIMessage(content=answer_str))
        
        # trim history if too large (keep last 20 messages)
        if len(chat_history) > 20:
            _sessions[session_id] = chat_history[-20:]
            
        return {"reply": answer_str, "rag_used": rag_used}
    except Exception as e:
        return _format_chat_rate_limit_error(e)


def _format_chat_rate_limit_error(e: Exception):
    """Parse Groq rate limit errors. Returns dict with wait_seconds if rate limited, string otherwise."""
    import re
    err_str = str(e).lower()
    print(f"[Chat Agent Error] {e}")  # Log for debugging

    # Only match REAL rate-limit errors — not generic "tokens" mentions
    rate_limit_keywords = [
        "rate_limit", "rate limit", "429",
        "tokens per minute", "tokens per day", "requests per minute",
        "requests per day", "token_limit", "quota", "limit exceeded",
        "resource_exhausted", "too many requests",
    ]

    if any(kw in err_str for kw in rate_limit_keywords):
        wait_seconds = None

        # Pattern: "try again in 1m23.456s" or "try again in 59.546s" or "try again in 2h1m30s"
        match = re.search(r'try again in\s+(?:(\d+)h)?\s*(?:(\d+)m)?\s*(?:(\d+(?:\.\d+)?)s)?', err_str)
        if match and any(match.group(i) for i in (1, 2, 3)):
            h = int(match.group(1) or 0)
            m = int(match.group(2) or 0)
            s = float(match.group(3) or 0)
            wait_seconds = h * 3600 + m * 60 + int(s) + (1 if s % 1 > 0 else 0)  # round up

        # Pattern: "retry after X seconds"
        if not wait_seconds:
            match = re.search(r'retry.?after\s+(\d+)', err_str)
            if match:
                wait_seconds = int(match.group(1))

        # Pattern: bare "N seconds"
        if not wait_seconds:
            match = re.search(r'(\d+)\s*seconds?', err_str)
            if match:
                wait_seconds = int(match.group(1))

        # Fallback: only if we truly couldn't parse anything
        if not wait_seconds:
            wait_seconds = 30

        return {"rate_limited": True, "wait_seconds": wait_seconds}

    return f"Agent error: {str(e)}"
