import os
import re
import pymupdf as fitz
import uuid
import tempfile
import faiss
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware  # ✨ Added CORS Support
from sentence_transformers import SentenceTransformer
from langchain_text_splitters import RecursiveCharacterTextSplitter
from openai import OpenAI
from graphviz import Digraph
from xml.sax.saxutils import escape

# ReportLab Layout Imports
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as RLImage, KeepTogether
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.pagesizes import letter
from reportlab.lib import colors
from reportlab.lib.utils import ImageReader

# =========================================================
# FASTAPI APP WITH CORS ENABLED
# =========================================================
app = FastAPI()

# Allow your Streamlit app frontend to communicate freely with the backend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# =========================================================
# LLM CONFIG
# =========================================================
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

client = OpenAI(
    api_key=GROQ_API_KEY,
    base_url="https://api.groq.com/openai/v1"
)

# =========================================================
# EMBEDDING MODEL & STORE GLOBALS
# =========================================================
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

global_index = None
global_chunks = []
global_summary = ""
global_flowchart_path = ""  # ✨ FIX: Initialized tracking global variable

# =========================================================
# HELPERS
# =========================================================
def extract_text_from_pdf(pdf_path):
    text = ""
    pdf_document = fitz.open(pdf_path)
    for page in pdf_document:
        text += page.get_text("text")
    return text

def create_chunks(text):
    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=120)
    return splitter.split_text(text)

def build_vector_store(chunks):
    global global_index, global_chunks
    global_chunks = chunks
    embeddings = embedding_model.encode(chunks,normalize_embeddings=True)
    embeddings = np.array(embeddings).astype("float32")
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatIP(dimension)
    index.add(embeddings)
    global_index = index

def search_chunks(query, top_k=4):
    query_embedding = embedding_model.encode([query],normalize_embeddings=True)
    query_embedding = np.array(query_embedding).astype("float32")
    distances, indices = global_index.search(query_embedding, top_k)
    return [global_chunks[idx] for idx in indices[0]]

# =========================================================
# GENERATE RESPONSES
# =========================================================
def generate_llm_response(prompt, temperature=0.15, max_tokens=1500):
    try:
        response = client.chat.completions.create(
            model="openai/gpt-oss-20b",
            messages=[
                {
                    "role": "system",
                    "content": """
You are a professional academic research paper analyzer.

Return ONLY the final answer.

STRICT RULES:
- Never show your reasoning.
- Never show a thinking process.
- Never write "Here's a thinking process".
- Never output HTML.
- Never output XML.
- Never output <br>, <br/>, <para>, </para>, <div>, <span>.
- Never output JSON.
- Never use code fences.
- Never create Markdown tables.
- Use clean Markdown headings.
- Use bullet points.
- Use short readable paragraphs.
- Do not invent information.
- Use only information supported by the research paper.
"""
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=temperature,
            max_tokens=max_tokens
        )

        result = response.choices[0].message.content

        if not result:
            return "No response generated by the model."

        return clean_llm_output(result)

    except Exception as e:
        return f"Error: {str(e)}"

def clean_llm_output(text: str) -> str:
    if not text:
        return ""

    import html

    # Decode escaped HTML entities
    text = html.unescape(text)

    # Remove ALL HTML/XML tags
    text = re.sub(
        r"<[^>]+>",
        "",
        text,
        flags=re.IGNORECASE
    )

    # Remove code fences
    text = re.sub(
        r"```(?:markdown|md|text)?",
        "",
        text,
        flags=re.IGNORECASE
    )
    text = text.replace("```", "")

    # Remove Markdown table separator lines
    text = re.sub(
        r"(?m)^\s*\|?[\s:-]+\|[\s|:-]*\|?\s*$",
        "",
        text
    )

    # Remove common reasoning prefixes
    text = re.sub(
        r"^(?:here(?:'s| is)\s+(?:a\s+)?thinking process:).*?\n",
        "",
        text,
        flags=re.IGNORECASE | re.DOTALL
    )

    # Remove accidental "thinking process" blocks
    text = re.sub(
        r"(?is)(?:^|\n)(?:thinking process|analysis|reasoning)\s*:.*?(?=\n(?:##|\*\*\d+\.))",
        "",
        text
    )

    # Remove common reasoning/thinking prefixes
    text = re.sub(
        r"(?is)^(?:here(?:'s| is)\s+)?(?:a\s+)?(?:thinking process|analysis|reasoning)\s*:.*?(?=##\s*\d+\.|\*\*\d+\.)",
        "",
        text
    )

    # Remove standalone thinking headings
    text = re.sub(
        r"(?im)^\s*(thinking process|analysis|reasoning)\s*:?\s*$",
        "",
        text
    )    
    
    # Fix escaped characters
    text = text.replace("\\n", "\n")
    text = text.replace("\\t", " ")

    # Normalize spaces
    text = re.sub(r"[ \t]+", " ", text)

    # Maximum 2 consecutive newlines
    text = re.sub(r"\n{3,}", "\n\n", text)

    return text.strip()
    
# =========================================================
# CREATE PDF REPORT
# =========================================================
def create_pdf_report(summary, flowchart_img_path=None):
    report_path = "research_report.pdf"
    
    doc = SimpleDocTemplate(
        report_path, 
        pagesize=letter,
        rightMargin=54, leftMargin=54, topMargin=54, bottomMargin=54
    )
    
    styles = getSampleStyleSheet()
    
    title_style = ParagraphStyle(
        'DocTitle',
        parent=styles['Title'],
        fontName='Helvetica-Bold',
        fontSize=24,
        leading=28,
        textColor=colors.HexColor('#1E293B'),
        alignment=0,
        spaceAfter=15
    )
    
    body_style = ParagraphStyle(
        'DocBody',
        parent=styles['BodyText'],
        fontName='Helvetica',
        fontSize=10.5,
        leading=16,
        textColor=colors.HexColor('#334155'),
        spaceAfter=12
    )
    
    heading_style = ParagraphStyle(
        'SectionHeading',
        parent=styles['Heading2'],
        fontName='Helvetica-Bold',
        fontSize=14,
        leading=18,
        textColor=colors.HexColor('#0F172A'),
        spaceBefore=14,
        spaceAfter=8,
        keepWithNext=True
    )

    subheading_style = ParagraphStyle(
        'SubSectionHeading',
        parent=styles['Heading3'],
        fontName='Helvetica-Bold',
        fontSize=12,
        leading=16,
        textColor=colors.HexColor('#334155'),
        spaceBefore=10,
        spaceAfter=6,
        keepWithNext=True
    )
    story = []
    
    story.append(Paragraph("Executive Research Analysis Report", title_style))
    story.append(Spacer(1, 10))
    
    safe_summary = clean_llm_output(summary)

    for block in safe_summary.split("\n\n"):

        block = block.strip()

        if not block:
            continue

        # Main headings
        if block.startswith("## "):
            heading = block.replace("## ", "", 1).strip()
            story.append(
                Paragraph(
                    escape(heading),
                    heading_style
                )
            )

        # Subheadings
        elif block.startswith("### "):
            subheading = block.replace("### ", "", 1).strip()

            story.append(
                Paragraph(
                    escape(subheading),
                    subheading_style
                )
            )

        # Bullet points
        elif any(line.strip().startswith("- ") for line in block.splitlines()):

            for line in block.splitlines():
                line = line.strip()

                if not line:
                    continue

                if line.startswith("- "):
                    bullet = escape(line[2:].strip())

                    story.append(
                        Paragraph(
                            "• " + bullet,
                            body_style
                        )
                    )
                else:
                    story.append(
                        Paragraph(
                            escape(line),
                            body_style
                        )
                    )

        # Normal paragraph
        else:
            safe_block = escape(block).replace("\n", " ")

            story.append(
                Paragraph(
                    safe_block,
                    body_style
                )
            )

        story.append(Spacer(1, 6))
    
     
    
    if flowchart_img_path and os.path.exists(flowchart_img_path):
        img = ImageReader(flowchart_img_path)
        orig_w, orig_h = img.getSize()
        
        MAX_WIDTH = 450
        MAX_HEIGHT = 550
        
        scale = min(MAX_WIDTH / float(orig_w), MAX_HEIGHT / float(orig_h))
        final_w = orig_w * scale
        final_h = orig_h * scale

        diagram_elements = [
            Spacer(1, 15),
            Paragraph("System Workflow & Logic Structural Diagram", heading_style),
            Spacer(1, 15),
            RLImage(flowchart_img_path, width=final_w, height=final_h)
        ]
        story.append(KeepTogether(diagram_elements))
        
    doc.build(story)
    return report_path

# =========================================================
# ENDPOINTS
# =========================================================
@app.post("/analyze-paper")
async def analyze_paper(pdf_file: UploadFile = File(None), text: str = Form(None)):
    global global_summary, global_flowchart_path
    text_data = ""
    try:
        if pdf_file:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
                temp_pdf.write(await pdf_file.read())
                temp_path = temp_pdf.name
            text_data = extract_text_from_pdf(temp_path)
        elif text:
            text_data = text
        else:
            return {"error": "No PDF or text data detected"}

        chunks = create_chunks(text_data)
        build_vector_store(chunks)
        analysis_context = "\n\n".join(chunks[:25])
        summary_prompt = f"""
        Analyze the research paper below and create a professional
        academic executive summary.

        IMPORTANT OUTPUT RULES:

        1. Return ONLY the final answer.
        2. Do NOT show reasoning.
        3. Do NOT show a thinking process.
        4. Do NOT output HTML or XML.
        5. Do NOT use <br>, <br/>, <para>, </para>, <div>, or <span>.
        6. Do NOT use JSON.
        7. Do NOT use code fences.
        8. Do NOT use Markdown tables.
        9. Use normal Markdown.
        10. Use headings, bullet points and short paragraphs.
        11. Do not invent information.
        12. Use only information supported by the paper.

        Use EXACTLY this structure:

        ## 1. MAIN RESEARCH OBJECTIVE

        Explain:
        - The main objective of the study.
        - The research problem.
        - The scope of the study.
        - Important geographical or industry context.

        ## 2. KEY METHODOLOGY & FINDINGS

        ### Methodology

        Explain:
        - Research methodology.
        - Data/source type.
        - Research approach.

        ### Key Findings

        Provide 4-6 important findings as bullet points.

        ### Important Evidence

        Provide important:
        - statistics
        - measurements
        - market figures
        - experimental results

        Only include information actually supported by the paper.

        ### Key Challenges

        List the major technical, social, regulatory,
        environmental and practical challenges.

        ## 3. BUSINESS & REAL-WORLD IMPLICATIONS

        ### Business Implications

        Explain how the findings affect:
        - companies
        - industries
        - markets
        - investors
        - technology providers

        ### Real-World Implications

        Explain effects on:
        - users
        - society
        - infrastructure
        - safety
        - environment
        - public policy

        ### Future Outlook

        Explain the future direction discussed by the authors.

        ## Executive Takeaway

        Provide a concise 3-5 sentence conclusion.

        RESEARCH PAPER:

        {analysis_context}
        """
        summary = generate_llm_response(summary_prompt, temperature=0.15, max_tokens=1800)
        global_summary = summary

        flow_text = " ".join(chunks[:8])
        flow_prompt = f"Extract exactly 4-5 core logical milestones from this research process sequence flow.\nRules:\n- One absolute short step name per line\n- No sequence numbers or explanation text\n- Max 4 words per line\n\nPaper Text:\n{flow_text}"
        steps_text = generate_llm_response(flow_prompt, temperature=0.2, max_tokens=500)
        steps = [line.strip() for line in steps_text.split("\n") if line.strip() and len(line.strip()) < 40][:5]

        if len(steps) < 2:
            steps = ["Document Ingestion", "Feature Extraction", "Neural Aggregation", "Validation Verification", "Output Compilation"]

        # 3. Create clean Executive Corporate Stylized Flowchart (COMPACT VERSION)
        dot = Digraph()
        
        # Reduced ranksep to 0.25 to bring vertical steps much closer together
        dot.attr(rankdir="TB", dpi="200", ranksep="0.25")
        
        # Reduced fontsize to 13 and margin to "0.2,0.1" to make boxes significantly shorter
        dot.attr('node', fontname='Helvetica', fontsize='13', shape='box', style='filled, rounded', 
                 color='#3B82F6', fillcolor='#F0F9FF', fontcolor='#1E3A8A', penwidth='1.2', margin="0.2,0.1")
        dot.attr('edge', fontname='Helvetica', fontsize='9', color='#64748B', penwidth='1.2', arrowsize='0.7')

        for i, step in enumerate(steps):
            dot.node(str(i), f" {step} ")
        for i in range(len(steps) - 1):
            dot.edge(str(i), str(i + 1))

        filename = f"flowchart_{uuid.uuid4().hex}"
        output_img = dot.render(filename, format="png", cleanup=True)
        global_flowchart_path = output_img

        create_pdf_report(summary, output_img)

        return {
            "summary": summary,
            "flowchart_url": f"/flowchart/{os.path.basename(output_img)}",
            "download_url": "/download-report",
            "analyzer_steps": [
                {"title": f"{i+1}. {step}", "desc": ""}
                for i, step in enumerate(steps)
            ]
        }
    except Exception as e:
        return {"error": str(e)}

@app.post("/chat-with-paper")
async def chat_with_paper(question: str = Form(...)):
    try:
        if global_index is None:
            return {"error": "Please analyze a paper first"}
        relevant_chunks = search_chunks(question)
        context = "\n\n".join(relevant_chunks)
        prompt = f"""
        You are an academic research assistant.

        Answer the user's question using ONLY the supplied research-paper context.

        OUTPUT RULES:

        - Give the direct answer first.
        - Use clean Markdown only.
        - Use short paragraphs.
        - Use bullet points when useful.
        - Use ## or ### headings only when necessary.
        - Do not use Markdown tables.
        - Do not use HTML.
        - Do not use XML.
        - Do not use JSON.
        - Do not use code fences.
        - Do not show reasoning.
        - Do not show thinking.
        - Do not mention internal analysis.
        - Do not invent information.
        - Do not repeat the question.

        If the answer is not available in the supplied context, say:

         "The provided paper context does not contain enough information to answer this."

        RESEARCH PAPER CONTEXT:

        {context}

        USER QUESTION:

        {question}

        Return only the final answer.
        """
        return {"question": question, "answer": generate_llm_response(prompt, temperature=0.2, max_tokens=900)}
    except Exception as e:
        return {"error": str(e)}

@app.get("/download-report")
async def download_report():
    return FileResponse("research_report.pdf", media_type="application/pdf", filename="research_report.pdf")

@app.post("/generate-flowchart")
async def generate_flowchart():
    try:
        global global_chunks
        if not global_chunks:
            return {"error": "Analyze PDF first in the primary engine tab."}

        text_data = " ".join(global_chunks[:10])
        
        prompt = f"""
        You are extracting the research workflow from an academic paper.

        Return EXACTLY 5 lines.

        Each line MUST follow this format:

        TITLE | DESCRIPTION

        Rules:
  
        - Exactly 5 lines.
        - One milestone per line.
        - Title must contain 2-5 words.
        - Description must contain 8-20 words.
        - Do not number the lines.
        - Do not use Markdown.
        - Do not use tables.
        - Do not use HTML.
        - Do not use XML.
        - Do not explain your reasoning.
        - Do not add introductory or concluding text.
        - Do not invent information.

        Example:

        Problem Definition | Identifies the main research problem and objectives addressed by the study.
        Literature Review | Reviews previous research, theories, datasets, and industry evidence relevant to the study.
        Methodology | Describes the methods, technologies, experiments, or analytical framework used by researchers.
        Results Analysis | Examines the major findings, measurements, comparisons, and evidence reported in the research.
        Conclusions | Summarizes the research contribution, limitations, practical implications, and future directions.

        Research Paper:

        {text_data}
        """
         

        
        
        steps_text = generate_llm_response(prompt, temperature=0.2, max_tokens=1000)
        lines = [line.strip() for line in steps_text.split("\n") if "|" in line]
        
        extracted_steps = []
        for idx, line in enumerate(lines[:5]):
            parts = line.split("|", 1)
            title = parts[0].strip()
            
            # ✨ FIX: Clean any leading digits (e.g., '1. ', '2. ') injected by the LLM
            title = re.sub(r'^\d+[\.\s\-–]+', '', title)
            
            desc = parts[1].strip()
            extracted_steps.append({
                "title": f"{idx+1}. {title}",
                "desc": desc
            })
        
        if len(extracted_steps) < 3:
            extracted_steps = [
                {"title": "1. Core Objective Alignment", "desc": "The primary problem statement and initial experimental hypotheses are established from document constraints."},
                {"title": "2. Data Assembly & Preprocessing", "desc": "Raw source matrices, baseline metrics, or study populations are clean-filtered for processing."},
                {"title": "3. Experimental Framework Execution", "desc": "The main technical methodology, core algorithms, or operational tests are systematically deployed."},
                {"title": "4. Performance Metrics Validation", "desc": "Outputs are analyzed against rigorous control benchmarks to verify scientific validation accuracy."},
                {"title": "5. Insight Synthesis & Conclusions", "desc": "Final qualitative and quantitative findings are compiled alongside real-world implementation scopes."}
            ]
            
        return {"steps": extracted_steps}
    except Exception as e:
        return {"error": str(e)}
        
@app.get("/flowchart/{filename}")
async def get_flowchart(filename: str):
    return FileResponse(filename, media_type="image/png")

@app.get("/")
async def root():
    return {"message": "Research Paper RAG API Active"}