from fastapi import FastAPI
from fastapi.responses import FileResponse
from pydantic import BaseModel
from typing import List
from transformers import pipeline
from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from datetime import datetime
import uuid

app = FastAPI()

# Load model once at startup
generator = pipeline("text-generation", model="distilgpt2")

SECTIONS = [
    "Abstract",
    "Introduction",
    "Literature Review",
    "Methodology",
    "System Design",
    "Implementation",
    "Evaluation",
    "Conclusion and Future Work"
]

class TaskItem(BaseModel):
    task: str

class ReportRequest(BaseModel):
    tasks: List[TaskItem]

    def build_overview(self) -> str:
        overview = ["This project summarizes a year of academic effort including planning, development, and evaluation.\n"]
        for item in self.tasks:
            overview.append(f"- {item.task}")
        overview.append("")
        return "\n".join(overview)

def generate_section(section_name: str, overview: str) -> str:
    prompt = f"""
You are an academic writer. Write a section titled '{section_name}' for a graduation project report based on:

{overview}

Avoid conclusions or references. Use formal language.
"""
    result = generator(prompt, max_length=512, num_return_sequences=1)
    return result[0]['generated_text']

def create_word_doc(sections_content, filename="Graduation_Project_Report.docx"):
    doc = Document()
    style = doc.styles['Normal']
    font = style.font
    font.name = 'Times New Roman'
    font.size = Pt(12)

    doc.add_heading("Graduation Project Report", 0).alignment = WD_ALIGN_PARAGRAPH.CENTER
    doc.add_paragraph("Author: Your Name")
    doc.add_paragraph("Supervisor: Supervisor Name")
    doc.add_paragraph("Department: Computer Science")
    doc.add_paragraph("University: Your University")
    doc.add_paragraph("Date: " + datetime.today().strftime("%B %d, %Y"))
    doc.add_page_break()

    for section_name, section_text in sections_content.items():
        doc.add_heading(section_name, level=1)
        for paragraph in section_text.split('\n'):
            line = paragraph.strip()
            if line:
                doc.add_paragraph(line)
        doc.add_page_break()

    doc.save(filename)
    return filename

@app.post("/generate-report")
def generate_report(req: ReportRequest):
    overview = req.build_overview()

    sections = {}
    for section in SECTIONS:
        sections[section] = generate_section(section, overview)

    unique_filename = f"report_{uuid.uuid4().hex[:8]}.docx"
    create_word_doc(sections, unique_filename)

    return {"download_url": f"/download/{unique_filename}"}

@app.get("/download/{filename}")
def download_file(filename: str):
    return FileResponse(path=filename, filename=filename, media_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document')

@app.get("/")
def root():
    return {"status": "Running"}