#!/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: Guide the user on *how to achieve* their programming tasks. "
            "Provide general explanations and pseudocode examples when appropriate. Avoid full detailed code unless necessary."
        )
    else:  # teacher mode default
        base += (
            "You are in TEACHER MODE: Provide detailed suggestions, structured explanations, and full code examples."
        )
    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)