import os
import json
import logging
import httpx
from enum import Enum
from typing import Dict, List, Optional, Any
from datetime import datetime

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
logger = logging.getLogger("assessment_report_processor")

# System startup confirmation
logger.info("Initializing Question Bank Training Processor [Version 3.0.0-LIGHT].")

DATA_ROOT = os.environ.get("TRAINED_DATA_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "trained_data"))
os.makedirs(DATA_ROOT, exist_ok=True)

# Ollama Configuration
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434/api/generate")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "gemma3:4b")
OLLAMA_TIMEOUT = float(os.environ.get("OLLAMA_TIMEOUT", 30.0))

class TrainingStatus(str, Enum):
    QUEUED = "queued"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

def _client_dir(client_id: int) -> str:
    return os.path.join(DATA_ROOT, str(client_id))

def _ensure_client_dir(client_id: int) -> str:
    path = _client_dir(client_id)
    os.makedirs(path, exist_ok=True)
    return path

def _classify_with_ollama(question_text: str, category_hint: str, options: str = "") -> Optional[str]:
    """Uses Ollama gemma3:4b to classify the question into a professional competency."""
    if not question_text:
        return None

    # Professional Competency Standardization Rules
    prompt = f"""
    Act as a Senior Academic Consultant and Competency Mapping Expert.
    Your task is to identify the single most accurate 'Professional Competency' for the question provided.

    STANDARDIZATION GUIDELINES:
    1. Use FORMAL, FULL NAMES (e.g., use 'Mathematics' instead of 'Math').
    2. Use BROAD PROFESSIONAL DOMAINS (e.g., use 'English Language' instead of 'Grammar' or 'English Grammar').
    3. Use TITLE CASE (e.g., 'General Knowledge' instead of 'general knowledge').
    4. NO ABBREVIATIONS (e.g., 'General Knowledge' instead of 'GK').
    5. Prioritize the Question Text and Options; use the Category Hint only for context.
    6. Return ONLY the competency name. No prefixes, no punctuation, no extra words.

    EXAMPLES OF PREFERRED TERMS:
    - Arithmetic/Algebra/Calculus -> Mathematics
    - Grammar/Vocabulary/Verbal -> English Language
    - Logic/Aptitude -> Logical Reasoning
    - Programming/Code -> Software Development
    - Physics/Chemistry/Biology -> [Subject Name]
    - Historical Events/Dates -> History
    - Geography/Maps -> Geography
    - Sports/Athletes -> General Knowledge

    Question Text: {question_text}
    Options: {options}
    Category Hint: {category_hint}

    Professional Competency:"""

    try:
        response = httpx.post(
            OLLAMA_URL,
            json={
                "model": OLLAMA_MODEL,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": 0.0,
                    "top_p": 0.9
                }
            },
            timeout=OLLAMA_TIMEOUT
        )
        if response.status_code == 200:
            result = response.json().get("response", "").strip()

            # Post-processing normalization
            result = result.split("\n")[0].strip()
            result = result.rstrip(".:!").strip()

            # Manual Mapping for absolute consistency
            normalization_map = {
                "math": "Mathematics",
                "maths": "Mathematics",
                "grammar": "English Language",
                "english grammar": "English Language",
                "gk": "General Knowledge",
                "logic": "Logical Reasoning",
                "aptitude": "Logical Reasoning",
                "historical awareness": "History",
                "historical": "History"
            }

            if result.lower() in normalization_map:
                return normalization_map[result.lower()]

            if len(result.split()) <= 3:
                result = result.title()

            return result if result else None
    except Exception as e:
        logger.warning(f"Ollama classification failed: {e}")

    return None

def _classify_question(q: dict) -> dict:
    """Extracts competency and other metadata from a question record using LLM with fallback."""
    qid = str(q.get("id", ""))
    cat = str(q.get("categoryName") or "General Knowledge").strip()
    text = str(q.get("questionText") or "").strip()
    opts = str(q.get("questionOption") or "").strip()

    # Primary logic: Professional LLM Classification (Now includes options)
    skill = _classify_with_ollama(text, cat, opts)

    # Secondary/Fallback logic
    if not skill:
        root_domain = cat.split(">")[0].split("-")[0].strip()
        if not root_domain or root_domain.lower() in ["input type", "test", "uncategorized", "default"]:
            skill = "General Knowledge"
        else:
            skill = root_domain.title()
            if skill.lower() == "gk":
                skill = "General Knowledge"

    return {
        "question_id": qid,
        "competency": skill,
        "category": cat,
        "question_text_preview": (text[:100] + "...") if len(text) > 100 else text,
        "question_type": q.get("questionType")
    }

def get_training_status(client_id):
    path = os.path.join(_client_dir(client_id), "status.json")
    if not os.path.exists(path): return None
    try:
        with open(path, "r") as f:
            return json.load(f)
    except json.JSONDecodeError:
        return {"status": TrainingStatus.FAILED.value, "error": "Corrupted status file"}

def mark_training_failed(client_id: int, error_message: str) -> dict:
    path = _ensure_client_dir(client_id)
    meta = {"status": TrainingStatus.FAILED.value, "error": error_message, "timestamp": datetime.now().isoformat()}
    with open(os.path.join(path, "status.json"), "w") as f:
        json.dump(meta, f, indent=4)
    return meta

def get_question_bank_result(client_id):
    path = os.path.join(_client_dir(client_id), "question_bank_result.json")
    if not os.path.exists(path): return None
    with open(path, "r") as f:
        return json.load(f)

def list_all_clients():
    """Lists all clients in the trained_data directory."""
    if not os.path.exists(DATA_ROOT):
        return []
    clients = []
    for client_id in os.listdir(DATA_ROOT):
        if os.path.isdir(os.path.join(DATA_ROOT, client_id)):
            try:
                status = get_training_status(int(client_id))
                clients.append({"client_id": int(client_id), "status": status.get("status") if status else "unknown"})
            except ValueError:
                continue
    return clients

def delete_client_data(client_id):
    path = _client_dir(client_id)
    if os.path.exists(path):
        # Check if processing
        status = get_training_status(client_id)
        if status and status.get("status") == TrainingStatus.PROCESSING.value:
            raise RuntimeError("Cannot delete while training is in progress.")

        import shutil
        shutil.rmtree(path)
        return True
    return False

def initialize_client_store(client_id: int) -> dict:
    """Initialize or reset the client's data store for question bank training."""
    path = _ensure_client_dir(client_id)

    # Clean up old files
    for tmp_file in ["accumulated_questions.json", "question_bank_result.json"]:
        tmp_path = os.path.join(path, tmp_file)
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

    meta = {"status": TrainingStatus.PROCESSING.value, "timestamp": datetime.now().isoformat()}
    with open(os.path.join(path, "status.json"), "w") as f:
        json.dump(meta, f, indent=4)
    return meta

def accumulate_question_bank(client_id: int, questions: List[dict]) -> dict:
    path = _ensure_client_dir(client_id)
    store_file = os.path.join(path, "accumulated_questions.json")
    existing = []
    if os.path.exists(store_file):
        with open(store_file, "r") as f:
            existing = json.load(f)

    # Deduplicate by ID
    dedup = {str(q.get("id", "")): q for q in existing if q.get("id")}
    for q in questions:
        if q.get("id"):
            dedup[str(q.get("id", ""))] = q

    existing = list(dedup.values())
    with open(store_file, "w") as f:
        json.dump(existing, f)
    return {"total_questions": len(existing)}

def finalize_question_bank(client_id: int) -> dict:
    """Classifies all questions and saves to question_bank_result.json. Includes checkpointing."""
    path = _ensure_client_dir(client_id)
    store_file = os.path.join(path, "accumulated_questions.json")
    result_path = os.path.join(path, "question_bank_result.json")

    if not os.path.exists(store_file):
        error_msg = "No accumulated questions found."
        mark_training_failed(client_id, error_msg)
        return {"error": error_msg}

    try:
        with open(store_file, "r") as f:
            question_bank = json.load(f)

        competency_summary = {}
        mapped_questions = {}

        # Load existing progress if available (checkpointing)
        if os.path.exists(result_path):
            try:
                with open(result_path, "r") as f:
                    existing_result = json.load(f)
                    mapped_questions = existing_result.get("questions", {})
                    competency_summary = existing_result.get("competency_summary", {})
            except json.JSONDecodeError:
                pass # Ignore corrupt checkpoint

        total = len(question_bank)
        save_interval = 50  # Save progress every 50 questions
        
        for i, q in enumerate(question_bank):
            qid = str(q.get("id", ""))
            
            # Skip if already processed in a previous interrupted run
            if qid and qid in mapped_questions:
                continue

            if (i + 1) % 10 == 0 or (i + 1) == total:
                logger.info(f"[Client {client_id}] Classifying questions: {i+1}/{total}...")

            classification = _classify_question(q)
            comp = classification["competency"]

            mapped_questions[qid] = classification
            competency_summary[comp] = competency_summary.get(comp, 0) + 1
            
            # Checkpoint save
            if (i + 1) % save_interval == 0:
                result = {
                    "client_id": client_id,
                    "total_questions": len(question_bank),
                    "trained_at": datetime.now().isoformat(),
                    "competency_summary": competency_summary,
                    "questions": mapped_questions
                }
                with open(result_path, "w") as f:
                    json.dump(result, f, indent=4)

        # Final save
        result = {
            "client_id": client_id,
            "total_questions": len(question_bank),
            "trained_at": datetime.now().isoformat(),
            "competency_summary": competency_summary,
            "questions": mapped_questions
        }

        with open(result_path, "w") as f:
            json.dump(result, f, indent=4)

        # Mark COMPLETED
        meta = {
            "status": TrainingStatus.COMPLETED.value,
            "total_questions": len(question_bank),
            "timestamp": datetime.now().isoformat()
        }
        with open(os.path.join(path, "status.json"), "w") as f:
            json.dump(meta, f, indent=4)

        logger.info(f"[Client {client_id}] Question bank finalized: {len(question_bank)} questions.")
        return result

    except Exception as e:
        logger.exception(f"[Client {client_id}] Fatal error during question bank finalization: {e}")
        mark_training_failed(client_id, f"Internal error during background finalization: {str(e)}")
        return {"error": str(e)}