#!/usr/bin/env python3
import os
import json
import logging
import re
from typing import Dict, List, Optional
from pathlib import Path
from flask import Flask, request, jsonify
from flask_cors import CORS
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from typing_extensions import TypedDict


# --- Type Definitions ---
class AssistantState(TypedDict):
    conversationSummary: str
    language: str 
    mode: str  # "teacher" or "student"

# --- Logging ---
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("code-assistant")

# --- Load environment ---
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    logger.error("GROQ_API_KEY not set in environment")
    raise RuntimeError("GROQ_API_KEY not set in environment")

# --- Flask app setup ---
BASE_DIR = Path(__file__).resolve().parent
static_folder = BASE_DIR / "static"

app = Flask(__name__, static_folder=str(static_folder), static_url_path="/static")
CORS(app)

# --- LLM setup ---
llm = ChatGroq(
    model=os.getenv("LLM_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct"),
    temperature=0.1,
    max_tokens=2048,
    api_key=GROQ_API_KEY,
)

# --- Helper functions ---

def detect_language_from_text(text: str) -> Optional[str]:
    if not text:
        return None
    lower = text.lower()
    known_languages = ["python", "javascript", "java", "c++", "c#", "go", "ruby", "php", "typescript", "swift"]
    lang_match = re.search(r'\b(in|using|for)\s+(' + '|'.join(known_languages) + r')\b', lower)
    if lang_match:
        return lang_match.group(2).capitalize()
    return None

def update_summary(chat_history: List[Dict[str, str]]) -> str:
    """
    Simple heuristic summary: last 6 messages concatenated.
    Replace with your own summarization chain if desired.
    """
    recent_msgs = chat_history[-6:]
    summary = " | ".join(f"{m['role']}: {m['content'][:50].replace('\n',' ')}" for m in recent_msgs)
    return summary

# def build_system_prompt(language: str, conversation_summary: str, mode: str) -> str:
#     """
#     Build system prompt dynamically based on mode.
#     """
#     base = f"You are a helpful programming assistant. Current language: {language}. Conversation summary: {conversation_summary}\n\n"
    
#     if mode == "student":
#         base += (
#             "You are in STUDENT MODE: Your goal is to engage the user in problem-solving and learning. "
#             "Do NOT give complete answers or full code. Instead:\n"
#             " - Ask guiding questions to make the user think.\n"
#             " - Give hints, small examples, or pseudocode to help the user discover the solution.\n"
#             " - Encourage step-by-step problem solving and curiosity.\n"
#             " - Only provide full solutions as a last resort if the user is completely stuck."
#         )
#     else:  # teacher mode default
#         base += (
#             "You are in TEACHER MODE: Provide detailed suggestions, structured explanations, and full code examples. "
#             "Explain reasoning clearly and comprehensively."
#         )
#     return base
def build_system_prompt(language: str, conversation_summary: str, mode: str) -> str:
    """
    Build system prompt dynamically based on mode, restricting to code/problem-solving only.
    """
    base = (
        f"You are a helpful programming assistant. "
        f"Your sole purpose is to assist with coding, programming, debugging, and technical problem solving. "
        f"Current language: {language}. "
        f"Conversation summary: {conversation_summary}\n\n"
    )

    if mode == "student":
        base += (
            "You are in STUDENT MODE.\n"
            "Your goal is to engage the user ONLY in programming, coding, and problem-solving tasks.\n"
            "STRICT RULES:\n"
            " - ❌ Do NOT give complete answers or full code unless the user is completely stuck.\n"
            " - ❌ Do NOT answer general knowledge, personal, or unrelated questions (e.g., names, trivia, history, etc.).\n"
            " - ❌ Politely refuse any out-of-context or non-programming queries by replying: "
            "\"I'm here only to help with programming or technical problem-solving questions.\"\n"
            " - ✅ ALWAYS guide the user through the problem-solving process instead of directly giving an answer.\n"
            " - ✅ Ask guiding questions to make the user think about coding problems.\n"
            " - ✅ Give hints, small examples, or pseudocode to help the user discover the solution.\n"
            " - ✅ Encourage step-by-step problem solving and curiosity.\n"
        )
    else:  # teacher mode default
        base += (
            "You are in TEACHER MODE.\n"
            "Your goal is to provide detailed explanations, structured reasoning, and complete code examples when needed.\n"
            "STRICT RULES:\n"
            " - ❌ Only answer questions related to programming, coding, or technical problem solving.\n"
            " - ❌ Politely refuse any unrelated, personal, or general knowledge questions by replying: "
            "\"I'm here only to help with programming or technical problem-solving questions.\"\n"
            " - ✅ Provide clear reasoning, best practices, and full working examples for programming tasks.\n"
        )

    return base

    return base
# --- Routes ---

@app.route("/", methods=["GET"])
def serve_frontend():
    try:
        return app.send_static_file("frontend.html")
    except Exception:
        return "<h3>frontend.html not found in static/ — please add your frontend.html there.</h3>", 404

@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json(force=True)
    chat_history = data.get("chat_history", [])
    assistant_state = data.get("assistant_state", {})

    conversation_summary = assistant_state.get("conversationSummary", "")
    language = assistant_state.get("language", "Python")
    mode = assistant_state.get("mode", "teacher").lower()
    if mode not in ("teacher", "student"):
        mode = "teacher"

    # Detect language from last user message
    last_user_msg = ""
    for msg in reversed(chat_history):
        if msg.get("role") == "user" and msg.get("content"):
            last_user_msg = msg["content"]
            break
    detected_lang = detect_language_from_text(last_user_msg)
    if detected_lang and detected_lang.lower() != language.lower():
        logger.info(f"Detected new language: {detected_lang}")
        language = detected_lang

    # Build system prompt based on mode
    system_prompt = build_system_prompt(language, conversation_summary, mode)
    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(chat_history)

    try:
        llm_response = llm.invoke(messages)
        assistant_reply = llm_response.content if hasattr(llm_response, "content") else str(llm_response)
    except Exception as e:
        logger.exception("LLM invocation failed")
        return jsonify({
            "assistant_reply": "Sorry, the assistant is currently unavailable. Please try again later.",
            "updated_state": {
                "conversationSummary": conversation_summary,
                "language": language,
                "mode": mode,
            },
            "chat_history": chat_history,
        }), 500

    # Append assistant reply to chat history
    chat_history.append({"role": "assistant", "content": assistant_reply})

    # Update conversation summary
    conversation_summary = update_summary(chat_history)

    return jsonify({
        "assistant_reply": assistant_reply,
        "updated_state": {
            "conversationSummary": conversation_summary,
            "language": language,
            "mode": mode,
        },
        "chat_history": chat_history,
    })

@app.route("/ping", methods=["GET"])
def ping():
    return jsonify({"status": "ok"})

if __name__ == "__main__":
    port = int(os.getenv("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=True)