import gradio as gr
import requests
import os
import json
import tempfile
from fpdf import FPDF
from docx import Document

# Backend URLs (should be set via Environment Variables in HF Spaces)
CPU_BACKEND_URL = os.getenv("CPU_BACKEND_URL", "http://localhost:7860")
GPU_BACKEND_URL = os.getenv("GPU_BACKEND_URL", "http://localhost:7861")
API_KEY = os.getenv("API_KEY", "")

def get_backend_url(mode):
    return GPU_BACKEND_URL if mode == "GPU (Faster)" else CPU_BACKEND_URL

def health_check(mode):
    url = get_backend_url(mode)
    try:
        resp = requests.get(f"{url}/health", timeout=5)
        return resp.json()
    except Exception as e:
        return {"error": str(e)}

def upload_files(mode, files):
    if not files:
        return "No files selected."
    
    if len(files) > 3:
        return "Error: Maximum 3 files allowed."
    
    url = get_backend_url(mode)
    headers = {"X-API-Key": API_KEY} if API_KEY else {}
    
    results = []
    for file in files:
        with open(file.name, "rb") as f:
            file_payload = {"file": (os.path.basename(file.name), f)}
            try:
                resp = requests.post(f"{url}/ingest", files=file_payload, headers=headers)
                results.append(f"{os.path.basename(file.name)}: {resp.json()}")
            except Exception as e:
                results.append(f"{os.path.basename(file.name)}: Error - {str(e)}")
                
    return "\n".join(results)

def generate_files(content):
    # 1. TXT
    txt_fd, txt_path = tempfile.mkstemp(suffix=".txt")
    with os.fdopen(txt_fd, 'w', encoding='utf-8') as f:
        f.write(content)
        
    # 2. DOCX
    doc = Document()
    doc.add_heading('Course Syllabus', 0)
    # Simple split by lines to add as paragraphs
    for line in content.split('\n'):
        doc.add_paragraph(line)
    
    docx_fd, docx_path = tempfile.mkstemp(suffix=".docx")
    os.close(docx_fd)
    doc.save(docx_path)
    
    # 3. PDF
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    # Replace unsupported characters for basic FPDF
    safe_content = content.encode('latin-1', 'replace').decode('latin-1')
    pdf.multi_cell(0, 10, txt=safe_content)
    
    pdf_fd, pdf_path = tempfile.mkstemp(suffix=".pdf")
    os.close(pdf_fd)
    pdf.output(pdf_path)
    
    return [txt_path, docx_path, pdf_path]

def generate_syllabus(mode, course_name, description, constraints):
    url = get_backend_url(mode)
    headers = {"X-API-Key": API_KEY} if API_KEY else {}
    payload = {
        "course_name": course_name,
        "description": description,
        "constraints": constraints
    }
    
    try:
        resp = requests.post(f"{url}/generate_syllabus", json=payload, headers=headers)
        data = resp.json()
        
        syllabus_text = ""
        if "syllabus" in data:
            syllabus_text = data["syllabus"]
        else:
            syllabus_text = json.dumps(data, indent=2)
            
        # Generate downloadable files
        files = generate_files(syllabus_text)
        
        return syllabus_text, files, None # Success
        
    except Exception as e:
        return f"Error: {str(e)}", [], f"Error details: {str(e)}"

# UI Layout
with gr.Blocks(title="Open Source Syllabus Engine") as demo:
    gr.Markdown("# 🎓 Open Source Syllabus Engine")
    gr.Markdown("Ingest your documents and generate a production-grade syllabus using RAG.")
    
    with gr.Row():
        mode_radio = gr.Radio(
            choices=["CPU (Free)", "GPU (Faster)"],
            value="CPU (Free)",
            label="Select Backend"
        )
        health_btn = gr.Button("Check Connection")
        health_output = gr.JSON(label="Backend Status")
    
    health_btn.click(health_check, inputs=[mode_radio], outputs=[health_output])
    
    with gr.Tab("1. Ingest Documents"):
        gr.Markdown("Upload up to 3 documents (PDF, DOCX, TXT) to provide context for the syllabus.")
        file_input = gr.File(label="Upload Documents", file_count="multiple")
        ingest_btn = gr.Button("Ingest Files")
        ingest_output = gr.Text(label="Ingestion Result")
        ingest_btn.click(upload_files, inputs=[mode_radio, file_input], outputs=[ingest_output])

    with gr.Tab("2. Generate Syllabus"):
        with gr.Row():
            course_name = gr.Textbox(label="Course Name", placeholder="e.g. Advanced Machine Learning")
            course_desc = gr.Textbox(label="Description", placeholder="Describe what the course covers...")
        constraints = gr.Textbox(label="Optional Constraints", placeholder="e.g. 10 weeks, include a final project")
        gen_btn = gr.Button("Generate Syllabus", variant="primary")
        
        with gr.Row():
            gen_output = gr.Markdown(label="Generated Syllabus Preview")
            with gr.Column():
                gen_files = gr.File(label="Download Formats (.txt, .docx, .pdf)")
                error_box = gr.Text(label="Debug Info", visible=True)
                
        gen_btn.click(
            generate_syllabus, 
            inputs=[mode_radio, course_name, course_desc, constraints], 
            outputs=[gen_output, gen_files, error_box]
        )

    gr.Markdown("---")
    gr.Markdown("### Ops Notes")
    gr.Markdown("- Ensure `CPU_BACKEND_URL` and `GPU_BACKEND_URL` are set correctly.")
    gr.Markdown("- Use the API Key if authentication is enabled on the backends.")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)