import os
import re
import math
import tempfile
from typing import Tuple

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

from pypdf import PdfReader
from docx import Document


# =========================
# Model Configuration
# =========================

MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")

MAX_INPUT_CHARS = 12000
MAX_NEW_TOKENS = 700


# =========================
# Load Local Small LLM
# =========================

print(f"Loading model: {MODEL_ID}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    device_map="cpu",
    low_cpu_mem_usage=True,
)

model.eval()

print("Model loaded successfully.")


# =========================
# File Reading Functions
# =========================

def read_pdf(file_path: str) -> str:
    text = ""
    reader = PdfReader(file_path)

    for i, page in enumerate(reader.pages):
        page_text = page.extract_text() or ""
        text += f"\n\n--- Page {i + 1} ---\n{page_text}"

    return text.strip()


def read_docx(file_path: str) -> str:
    doc = Document(file_path)
    paragraphs = []

    for para in doc.paragraphs:
        if para.text.strip():
            paragraphs.append(para.text.strip())

    return "\n".join(paragraphs).strip()


def read_txt(file_path: str) -> str:
    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
        return f.read().strip()


def extract_text_from_file(file) -> str:
    if file is None:
        return ""

    file_path = file.name
    ext = os.path.splitext(file_path)[1].lower()

    try:
        if ext == ".pdf":
            return read_pdf(file_path)

        elif ext == ".docx":
            return read_docx(file_path)

        elif ext in [".txt", ".md"]:
            return read_txt(file_path)

        else:
            return "Unsupported file format. Please upload PDF, DOCX, TXT, or MD."
    except Exception as e:
        return f"Error reading file: {str(e)}"


# =========================
# Text Utility
# =========================

def clean_text(text: str) -> str:
    text = text.replace("\x00", " ")
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[ \t]{2,}", " ", text)
    return text.strip()


def limit_text(text: str, max_chars: int = MAX_INPUT_CHARS) -> str:
    text = clean_text(text)

    if len(text) <= max_chars:
        return text

    beginning = text[: int(max_chars * 0.65)]
    ending = text[-int(max_chars * 0.35):]

    return (
        beginning
        + "\n\n[... middle part shortened because report is long ...]\n\n"
        + ending
    )


def estimate_confidence(report_text: str) -> Tuple[int, str]:
    words = report_text.split()
    word_count = len(words)

    score = 50

    if word_count > 500:
        score += 10
    if word_count > 1200:
        score += 10
    if re.search(r"\bmethodology\b|\bmethod\b", report_text, re.I):
        score += 8
    if re.search(r"\bresult\b|\bresults\b|\banalysis\b", report_text, re.I):
        score += 8
    if re.search(r"\bconclusion\b", report_text, re.I):
        score += 6
    if re.search(r"\breference\b|\breferences\b|\bcitation\b", report_text, re.I):
        score += 5
    if re.search(r"\bobjective\b|\bgoal\b|\baim\b", report_text, re.I):
        score += 5

    if word_count < 250:
        score -= 15

    score = max(20, min(95, score))

    if score >= 80:
        level = "Strong defense readiness"
    elif score >= 65:
        level = "Moderate defense readiness"
    elif score >= 50:
        level = "Basic defense readiness"
    else:
        level = "Needs more preparation"

    return score, level


# =========================
# Local LLM Function
# =========================

def ask_local_llm(prompt: str) -> str:
    messages = [
        {
            "role": "system",
            "content": (
                "You are Assignment Defense AI. "
                "Your job is to help students defend their assignments in front of teachers. "
                "Use simple English. If useful, explain in Bangla-English style. "
                "Be practical, clear, and exam-focused. "
                "Do not invent facts that are not in the report."
            ),
        },
        {
            "role": "user",
            "content": prompt,
        },
    ]

    try:
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception:
        text = (
            "System: You are Assignment Defense AI.\n"
            f"User: {prompt}\n"
            "Assistant:"
        )

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=4096,
    )

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            temperature=0.4,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.12,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]

    response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    )

    return response.strip()


# =========================
# Main AI Features
# =========================

def generate_defense_pack(file, manual_text: str, difficulty: str, language_style: str):
    uploaded_text = extract_text_from_file(file)
    manual_text = manual_text.strip() if manual_text else ""

    if uploaded_text and manual_text:
        report_text = uploaded_text + "\n\nAdditional Notes:\n" + manual_text
    elif uploaded_text:
        report_text = uploaded_text
    elif manual_text:
        report_text = manual_text
    else:
        return (
            "Please upload a report or paste your assignment text.",
            "",
            "",
            "",
            "",
            "",
        )

    report_text = limit_text(report_text)

    confidence_score, confidence_level = estimate_confidence(report_text)

    prompt = f"""
You are preparing a student for assignment defense/viva.

Language style: {language_style}
Difficulty level: {difficulty}

Here is the student's report:

{report_text}

Create a complete defense preparation pack with these sections:

1. Report Summary
- Explain the assignment in simple words.
- Mention the main topic, objective, method, and result if available.

2. 12 Viva Questions
- Questions should be realistic and teacher-like.
- Include easy, medium, and challenging questions.

3. Easy Answers
- Give short, clear answers for each viva question.
- Answers should sound natural for a student.

4. Defense Speech
- Write a 1 to 2 minute formal speech for presenting this assignment.
- Make it confident but not overdramatic.

5. Weak Points Teacher May Ask
- Identify possible weak areas in the report.
- Give safe answer strategy for each weak point.

6. Final Preparation Tips
- Give practical tips before viva.
"""

    ai_output = ask_local_llm(prompt)

    confidence_report = f"""
# Confidence Score

**Score:** {confidence_score}/100  
**Level:** {confidence_level}

## Meaning
This score is estimated from your report structure, length, and presence of important academic sections.

## How to improve
- Understand the objective clearly.
- Memorize the methodology, not the full report.
- Prepare 5–10 key terms from your assignment.
- Practice explaining the report in 60 seconds.
- Be honest if you do not know an answer.
"""

    summary = ai_output
    viva_questions = extract_section(ai_output, "12 Viva Questions", "Easy Answers")
    easy_answers = extract_section(ai_output, "Easy Answers", "Defense Speech")
    speech = extract_section(ai_output, "Defense Speech", "Weak Points")
    weak_points = extract_section(ai_output, "Weak Points", "Final Preparation")

    return (
        summary,
        viva_questions,
        easy_answers,
        speech,
        weak_points,
        confidence_report,
    )


def extract_section(text: str, start_keyword: str, end_keyword: str) -> str:
    try:
        pattern = rf"(?is){re.escape(start_keyword)}(.*?){re.escape(end_keyword)}"
        match = re.search(pattern, text)

        if match:
            return match.group(1).strip()

        return "Section generated inside the full defense pack. Please check the Full Defense Pack tab."
    except Exception:
        return "Section extraction failed. Please check the Full Defense Pack tab."


def evaluate_practice_answer(question: str, student_answer: str, report_context: str):
    if not question.strip() or not student_answer.strip():
        return "Please provide both the viva question and your answer."

    report_context = limit_text(report_context or "", 4000)

    prompt = f"""
You are a strict but helpful viva teacher.

Report context:
{report_context}

Viva question:
{question}

Student answer:
{student_answer}

Evaluate the answer.

Give:
1. Score out of 10
2. What was good
3. What was weak
4. Better answer
5. One short tip for the student

Use simple English.
"""

    return ask_local_llm(prompt)


def create_custom_questions(file, manual_text: str, topic_focus: str, number_of_questions: int):
    uploaded_text = extract_text_from_file(file)
    manual_text = manual_text.strip() if manual_text else ""

    report_text = uploaded_text + "\n\n" + manual_text
    report_text = limit_text(report_text)

    if not report_text.strip():
        return "Please upload or paste your report first."

    prompt = f"""
Create {number_of_questions} viva questions from this report.

Focus area: {topic_focus if topic_focus else "overall assignment"}

For each question, include:
- Question
- Short easy answer
- Difficulty: Easy/Medium/Hard

Report:
{report_text}
"""

    return ask_local_llm(prompt)


# =========================
# Gradio UI
# =========================

custom_css = """
.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}
.main-title {
    text-align: center;
    font-size: 34px;
    font-weight: 800;
    margin-bottom: 8px;
}
.sub-title {
    text-align: center;
    font-size: 16px;
    opacity: 0.85;
    margin-bottom: 25px;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div class="main-title">Assignment Defense AI</div>
        <div class="sub-title">
        Upload your report, get viva questions, easy answers, defense speech, and confidence score.
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="Upload Report",
                file_types=[".pdf", ".docx", ".txt", ".md"],
            )

            manual_text = gr.Textbox(
                label="Or paste assignment text / notes",
                placeholder="Paste your assignment text here...",
                lines=10,
            )

            difficulty = gr.Radio(
                choices=["Easy", "Medium", "Hard"],
                value="Medium",
                label="Viva Difficulty",
            )

            language_style = gr.Radio(
                choices=[
                    "Simple English",
                    "Bangla-English Mixed",
                    "Formal Academic English",
                ],
                value="Simple English",
                label="Answer Style",
            )

            generate_btn = gr.Button(
                "Generate Defense Pack",
                variant="primary",
            )

        with gr.Column(scale=2):
            with gr.Tab("Full Defense Pack"):
                full_output = gr.Markdown()

            with gr.Tab("Viva Questions"):
                questions_output = gr.Markdown()

            with gr.Tab("Easy Answers"):
                answers_output = gr.Markdown()

            with gr.Tab("Defense Speech"):
                speech_output = gr.Markdown()

            with gr.Tab("Weak Points"):
                weak_output = gr.Markdown()

            with gr.Tab("Confidence Score"):
                confidence_output = gr.Markdown()

    generate_btn.click(
        fn=generate_defense_pack,
        inputs=[file_input, manual_text, difficulty, language_style],
        outputs=[
            full_output,
            questions_output,
            answers_output,
            speech_output,
            weak_output,
            confidence_output,
        ],
    )

    gr.Markdown("---")

    gr.Markdown("## Practice Viva Evaluator")

    with gr.Row():
        with gr.Column():
            practice_question = gr.Textbox(
                label="Viva Question",
                placeholder="Example: Why did you choose this methodology?",
                lines=3,
            )

            practice_answer = gr.Textbox(
                label="Your Answer",
                placeholder="Type your answer here...",
                lines=6,
            )

            report_context = gr.Textbox(
                label="Optional Report Context",
                placeholder="Paste a small part of your report if needed...",
                lines=6,
            )

            evaluate_btn = gr.Button("Evaluate My Answer")

        with gr.Column():
            evaluation_output = gr.Markdown()

    evaluate_btn.click(
        fn=evaluate_practice_answer,
        inputs=[practice_question, practice_answer, report_context],
        outputs=evaluation_output,
    )

    gr.Markdown("---")

    gr.Markdown("## Custom Viva Question Generator")

    with gr.Row():
        with gr.Column():
            topic_focus = gr.Textbox(
                label="Focus Topic",
                placeholder="Example: methodology, result, networking, algorithm, water cycle...",
            )

            number_of_questions = gr.Slider(
                minimum=5,
                maximum=25,
                value=10,
                step=1,
                label="Number of Questions",
            )

            custom_btn = gr.Button("Generate Custom Questions")

        with gr.Column():
            custom_output = gr.Markdown()

    custom_btn.click(
        fn=create_custom_questions,
        inputs=[file_input, manual_text, topic_focus, number_of_questions],
        outputs=custom_output,
    )

    gr.Markdown(
        """
        ### Notes
        - This app uses a small local Hugging Face model.
        - No API key is required.
        - First run may take time because the model downloads automatically.
        - For better quality, later you can add Gemini/Groq API as an optional mode.
        """
    )


if __name__ == "__main__":
    demo.launch()