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

# --- Type Definitions for State Management ---
class TaggedReply(TypedDict):
    reply: str
    tags: List[str]

class AssistantState(TypedDict):
    conversationSummary: str
    lastUserMessage: str
    language: str # New field to track the programming language
    taggedReplies: List[TaggedReply] # New field for saving/bookmarking replies

# --- 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")
    exit(1)

# --- 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 ---
# Using a model that's good for coding tasks
llm = ChatGroq(
    model=os.getenv("LLM_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct"), # Use the supported model
    temperature=0.1, 
    max_tokens=2048,
    api_key=GROQ_API_KEY,
)

PROGRAMMING_ASSISTANT_PROMPT = """
You are an expert programming assistant. Your role is to provide code suggestions, fix bugs, explain programming concepts, and offer contextual help based on the user's query and preferred programming language.

Behavior rules (follow these strictly):
- Contextual Help: Always aim to provide the most helpful, clear, and accurate information.
- Code Suggestions: When suggesting code, always enclose it in appropriate markdown code blocks (e.g., ```python\n...\n```).
- Error Explanation: When an error is provided, explain the root cause and provide a corrected code snippet if possible.
- Conceptual Questions: For questions like "What is a loop?", provide a clear, concise explanation with a simple, illustrative code example in the user's current language (if known, otherwise Python/JavaScript).
- Language Adaptation: Adjust your suggestions, code, and explanations to the programming language specified in the 'language' field of the 'AssistantState'. If 'language' is not set, ask the user what language they are working in.

STRICT OUTPUT FORMAT (JSON ONLY):
Return a single JSON object with the following keys:
- assistant_reply: string  // a natural language reply to the user (short, helpful, always present)
- state_updates: object  // updates to the internal state, keys may include: language, conversationSummary
- suggested_tags: array of strings // a list of 1-3 relevant tags for the assistant_reply (e.g., "Python", "Debugging", "Loop Concept")

Rules:
- ALWAYS include `assistant_reply` as a non-empty string.
- Do NOT produce any text outside the JSON object.
- Be concise in `assistant_reply`, but ensure the information is complete.
- Do not make up information.
"""

def extract_json_from_llm_response(raw_response: str) -> dict:
    default = {
        "assistant_reply": "I'm sorry — I couldn't understand that. Could you please rephrase?",
        "state_updates": {},
        "suggested_tags": [],
    }
    if not raw_response or not isinstance(raw_response, str):
        return default
    m = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", raw_response)
    json_string = m.group(1).strip() if m else raw_response
    first = json_string.find('{')
    last = json_string.rfind('}')
    
    if first != -1 and last != -1 and first < last:
        candidate = json_string[first:last+1]
    else:
        candidate = json_string # Fallback to the whole string if braces aren't clear

    candidate = re.sub(r',\s*(?=[}\]])', '', candidate)
    
    try:
        parsed = json.loads(candidate)
    except Exception as e:
        logger.warning("Failed to parse JSON from LLM output: %s. Raw candidate: %s", e, candidate)
        return default
    
    if isinstance(parsed, dict) and "assistant_reply" in parsed and isinstance(parsed["assistant_reply"], str) and parsed["assistant_reply"].strip():
        parsed.setdefault("state_updates", {})
        parsed.setdefault("suggested_tags", [])
        return parsed
    else:
        logger.warning("Parsed JSON missing 'assistant_reply' or invalid format. Returning default.")
        return default

# --- Flask 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)
    if not isinstance(data, dict):
        return jsonify({"error": "invalid request body"}), 400

    chat_history: List[Dict[str, str]] = data.get("chat_history") or []
    assistant_state: AssistantState = data.get("assistant_state") or {}

    state: AssistantState = {
        "conversationSummary": assistant_state.get("conversationSummary", ""),
        "lastUserMessage": "",
        "language": assistant_state.get("language", "Python"),
        "taggedReplies": assistant_state.get("taggedReplies", []),
    }
    
    for msg in reversed(chat_history):
        if msg.get("role") == "user" and msg.get("content"):
            state["lastUserMessage"] = msg["content"]
            break

    last_msg_lower = state["lastUserMessage"].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', last_msg_lower)
    if lang_match:
        detected_lang = lang_match.group(2).capitalize()
        if detected_lang.lower() != state["language"].lower():
            logger.info("Detected new language: %s", detected_lang)
            state["language"] = detected_lang
    
    action_hint = ""
    if state["language"]:
        action_hint = f"Focus your answer on the {state['language']} programming language. If the user asks a conceptual question, use {state['language']} for examples."
    else:
        action_hint = "The current language is unknown. Please ask the user to specify the programming language they are working in."

    user_prompt = f"""
Current State: {json.dumps({"language": state["language"], "summary": state["conversationSummary"][:200]})}
Last user message: {state["lastUserMessage"]}

SYSTEM_HINT: {action_hint}

Return ONLY valid JSON with keys: assistant_reply, state_updates, suggested_tags.
"""

    messages = [
        {"role": "system", "content": PROGRAMMING_ASSISTANT_PROMPT},
        {"role": "user", "content": user_prompt}
    ]

    try:
        logger.info("Invoking LLM for code assistant...")
        llm_response = llm.invoke(messages)
        raw_response = llm_response.content if hasattr(llm_response, "content") else str(llm_response)

        logger.info(f"Raw LLM response: {raw_response[:200]}...")
        parsed_result = extract_json_from_llm_response(raw_response)

    except Exception as e:
        logger.exception("LLM invocation failed")
        return jsonify({"error": "LLM invocation failed", "detail": str(e)}), 500

    updated_state_from_llm = parsed_result.get("state_updates", {})
    
    if 'conversationSummary' in updated_state_from_llm:
        state["conversationSummary"] = updated_state_from_llm["conversationSummary"]
    if 'language' in updated_state_from_llm:
        state["language"] = updated_state_from_llm["language"]

    assistant_reply = parsed_result.get("assistant_reply")
    if not assistant_reply or not isinstance(assistant_reply, str) or not assistant_reply.strip():
        assistant_reply = "I'm here to help with your code! What programming language are you using?"

    response_payload = {
        "assistant_reply": assistant_reply,
        "updated_state": state,
        "suggested_tags": parsed_result.get("suggested_tags", []),
    }

    return jsonify(response_payload)

@app.route("/tag_reply", methods=["POST"])
def tag_reply():
    data = request.get_json(force=True)
    if not isinstance(data, dict):
        return jsonify({"error": "invalid request body"}), 400

    reply_content = data.get("reply")
    tags = data.get("tags")
    assistant_state: AssistantState = data.get("assistant_state") or {}

    if not reply_content or not tags:
        return jsonify({"error": "Missing 'reply' or 'tags' in request"}), 400
    
    tags = [str(t).strip() for t in tags if str(t).strip()]
    if not tags:
        return jsonify({"error": "Tags list cannot be empty"}), 400

    state: AssistantState = {
        "conversationSummary": assistant_state.get("conversationSummary", ""),
        "lastUserMessage": "",
        "language": assistant_state.get("language", "Python"),
        "taggedReplies": assistant_state.get("taggedReplies", []),
    }

    new_tagged_reply: TaggedReply = {
        "reply": reply_content,
        "tags": tags,
    }

    state["taggedReplies"].append(new_tagged_reply)
    
    logger.info("Reply tagged with: %s", tags)

    return jsonify({
        "message": "Reply saved and tagged successfully.",
        "updated_state": state,
    }), 200

@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)