import gradio as gr
from groq import Groq
from fpdf import FPDF
import json
import os
import re
from datetime import datetime

# Initialize Groq client from environment variable
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY not found in environment variables. Please set it in Space Secrets.")

client = Groq(api_key=GROQ_API_KEY)

def clean_json_string(content):
    """Remove invalid control characters from JSON string"""
    content = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]', '', content)
    return content

def generate_roadmap(domain, level, time_available, time_unit):
    """Generate roadmap using Groq API"""
    if time_unit == "Weeks":
        total_weeks = int(time_available)
    elif time_unit == "Months":
        total_weeks = int(time_available) * 4
    elif time_unit == "Years":
        total_weeks = int(time_available) * 52
    
    system_prompt = """You are an expert educational curriculum designer. Return only valid JSON without any markdown formatting or explanatory text."""
    
    user_prompt = f"""Create a detailed learning roadmap for {domain} at {level} level.
Duration: {time_available} {time_unit.lower()} ({total_weeks} weeks).

STRICT REQUIREMENTS:
1. Return ONLY valid JSON, no markdown, no backticks, no explanation
2. Use this exact structure:
{{
    "domain": "{domain}",
    "level": "{level}",
    "duration": "{time_available} {time_unit}",
    "overview": "Brief overview text here",
    "phases": [
        {{
            "phase_name": "Phase Name",
            "duration_weeks": number,
            "description": "Description here",
            "topics": [
                {{
                    "name": "Topic name",
                    "resources": ["resource 1", "resource 2"],
                    "practice_project": "Project description"
                }}
            ],
            "milestones": ["milestone 1", "milestone 2"]
        }}
    ],
    "essential_resources": {{
        "courses": ["course 1"],
        "books": ["book 1"],
        "communities": ["community 1"],
        "tools": ["tool 1"]
    }},
    "tips_for_success": ["tip 1", "tip 2"],
    "next_steps": "What to do next"
}}

3. Ensure all strings are properly escaped (no raw newlines or tabs in strings)
4. Use \\n for newlines within strings if needed"""

    try:
        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.7,
            max_tokens=4096
        )
        
        content = response.choices[0].message.content.strip()
        
        # Clean up markdown code blocks
        if content.startswith("```json"):
            content = content[7:]
        elif content.startswith("```"):
            content = content[3:]
        if content.endswith("```"):
            content = content[:-3]
        
        content = content.strip()
        content = clean_json_string(content)
        
        # Try to parse JSON
        try:
            data = json.loads(content)
            return data
        except json.JSONDecodeError as e:
            # Try to fix common issues
            content = re.sub(r'(".*?)\n(.*?")', r'\1\\n\2', content, flags=re.DOTALL)
            content = re.sub(r'(".*?)\t(.*?")', r'\1\\t\2', content, flags=re.DOTALL)
            data = json.loads(content)
            return data
            
    except Exception as e:
        return {"error": f"Failed to generate roadmap: {str(e)}"}

def format_roadmark_for_display(roadmap_data):
    """Format roadmap as HTML"""
    if "error" in roadmap_data:
        return f"<p style='color:red; padding:20px;'>{roadmap_data['error']}</p>", None
    
    try:
        html_output = f"""
        <div style="font-family: 'Segoe UI', sans-serif; max-width: 900px; margin: 0 auto;">
            <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 30px; border-radius: 15px; margin-bottom: 20px;">
                <h1 style="margin: 0;">🎯 {roadmap_data.get('domain', 'Unknown')} Roadmap</h1>
                <p>Level: {roadmap_data.get('level', 'N/A')} | Duration: {roadmap_data.get('duration', 'N/A')}</p>
            </div>
            <div style="background: #f8f9fa; padding: 20px; border-radius: 10px; margin-bottom: 20px;">
                <h3>Overview</h3>
                <p>{roadmap_data.get('overview', 'No overview provided').replace(chr(10), '<br>')}</p>
            </div>
        """
        
        for i, phase in enumerate(roadmap_data.get('phases', []), 1):
            html_output += f"""
            <div style="border: 2px solid #e9ecef; border-radius: 15px; margin-bottom: 20px; overflow: hidden;">
                <div style="background: #667eea; color: white; padding: 15px;">
                    <h3 style="margin:0;">Phase {i}: {phase.get('phase_name', 'Unnamed')}</h3>
                    <p style="margin:5px 0 0 0;">Duration: {phase.get('duration_weeks', 'N/A')} weeks</p>
                </div>
                <div style="padding: 20px;">
                    <p>{phase.get('description', '').replace(chr(10), '<br>')}</p>
                    <h4>Topics:</h4>
            """
            
            for topic in phase.get('topics', []):
                html_output += f"<div style='background:#f8f9fa; padding:10px; margin:10px 0; border-radius:5px;'>"
                html_output += f"<p style='margin:0;'><b>{topic.get('name', 'Unnamed')}</b></p>"
                html_output += f"<ul style='margin:5px 0;'>"
                for res in topic.get('resources', []):
                    html_output += f"<li>{res}</li>"
                html_output += f"</ul>"
                html_output += f"<p style='background:#fff3cd; padding:5px; margin:5px 0; border-radius:3px;'><b>Project:</b> {topic.get('practice_project', 'N/A')}</p>"
                html_output += f"</div>"
            
            if phase.get('milestones'):
                html_output += f"<h4>Milestones:</h4><ul>"
                for ms in phase['milestones']:
                    html_output += f"<li>☐ {ms}</li>"
                html_output += f"</ul>"
            
            html_output += "</div></div>"
        
        html_output += "</div>"
        return html_output, json.dumps(roadmap_data)
        
    except Exception as e:
        return f"<p style='color:red; padding:20px;'>Error formatting display: {str(e)}</p>", None

def generate_pdf(roadmap_data, domain):
    """Generate PDF file"""
    try:
        if isinstance(roadmap_data, str):
            roadmap_data = json.loads(roadmap_data)
        
        if not roadmap_data or "error" in roadmap_data:
            return None
        
        pdf = FPDF()
        pdf.set_auto_page_break(auto=True, margin=15)
        
        def clean(text):
            if text is None:
                return ""
            text = str(text)
            return text.encode('latin-1', 'replace').decode('latin-1')
        
        # Title Page
        pdf.add_page()
        pdf.set_font("Arial", 'B', 24)
        pdf.set_text_color(102, 126, 234)
        pdf.cell(0, 20, clean(f"Learning Roadmap: {roadmap_data.get('domain', domain)}"), ln=True, align='C')
        
        pdf.set_font("Arial", '', 12)
        pdf.set_text_color(100, 100, 100)
        pdf.cell(0, 10, clean(f"Level: {roadmap_data.get('level', 'N/A')} | Duration: {roadmap_data.get('duration', 'N/A')}"), ln=True, align='C')
        pdf.ln(10)
        
        # Overview
        pdf.set_font("Arial", 'B', 16)
        pdf.set_text_color(0, 0, 0)
        pdf.cell(0, 10, "Overview", ln=True)
        pdf.set_font("Arial", '', 11)
        pdf.multi_cell(0, 6, clean(roadmap_data.get('overview', 'No overview')))
        pdf.ln(5)
        
        # Phases
        for i, phase in enumerate(roadmap_data.get('phases', []), 1):
            pdf.add_page()
            pdf.set_font("Arial", 'B', 18)
            pdf.set_fill_color(102, 126, 234)
            pdf.set_text_color(255, 255, 255)
            pdf.cell(0, 12, clean(f"Phase {i}: {phase.get('phase_name', f'Phase {i}')}"), ln=True, fill=True)
            
            pdf.set_text_color(100, 100, 100)
            pdf.set_font("Arial", 'I', 11)
            pdf.cell(0, 8, f"Duration: {phase.get('duration_weeks', 'N/A')} weeks", ln=True)
            pdf.ln(2)
            
            pdf.set_text_color(0, 0, 0)
            pdf.set_font("Arial", '', 11)
            pdf.multi_cell(0, 6, clean(phase.get('description', '')))
            pdf.ln(5)
            
            pdf.set_font("Arial", 'B', 13)
            pdf.set_text_color(102, 126, 234)
            pdf.cell(0, 8, "Topics & Resources:", ln=True)
            
            for topic in phase.get('topics', []):
                pdf.set_font("Arial", 'B', 11)
                pdf.set_text_color(0, 0, 0)
                pdf.cell(0, 6, clean(f"- {topic.get('name', 'Topic')}"), ln=True)
                
                pdf.set_font("Arial", '', 10)
                pdf.set_text_color(80, 80, 80)
                for resource in topic.get('resources', []):
                    pdf.cell(10)
                    pdf.cell(0, 5, clean(f"  * {resource}"), ln=True)
                
                pdf.set_text_color(200, 100, 0)
                pdf.set_font("Arial", 'B', 10)
                pdf.cell(0, 5, clean(f"Project: {topic.get('practice_project', 'N/A')}"), ln=True)
                pdf.ln(3)
            
            if phase.get('milestones'):
                pdf.set_font("Arial", 'B', 12)
                pdf.set_text_color(40, 167, 69)
                pdf.cell(0, 8, "Milestones:", ln=True)
                pdf.set_font("Arial", '', 10)
                pdf.set_text_color(0, 0, 0)
                for milestone in phase.get('milestones', []):
                    pdf.cell(10)
                    pdf.cell(0, 5, clean(f"[ ] {milestone}"), ln=True)
                pdf.ln(5)
        
        # Resources
        if roadmap_data.get('essential_resources'):
            pdf.add_page()
            pdf.set_font("Arial", 'B', 18)
            pdf.set_fill_color(102, 126, 234)
            pdf.set_text_color(255, 255, 255)
            pdf.cell(0, 12, "Essential Resources", ln=True, fill=True)
            pdf.ln(5)
            
            resources = roadmap_data['essential_resources']
            sections = [
                ("Recommended Courses", resources.get('courses', [])), 
                ("Books", resources.get('books', [])),
                ("Communities", resources.get('communities', [])),
                ("Tools", resources.get('tools', []))
            ]
            
            for title, items in sections:
                if items:
                    pdf.set_font("Arial", 'B', 13)
                    pdf.set_text_color(102, 126, 234)
                    pdf.cell(0, 8, clean(title), ln=True)
                    pdf.set_font("Arial", '', 10)
                    pdf.set_text_color(0, 0, 0)
                    for item in items:
                        pdf.cell(5)
                        pdf.cell(0, 5, clean(f"- {item}"), ln=True)
                    pdf.ln(3)
        
        # Save to /tmp for Hugging Face
        safe_domain = "".join(c for c in domain if c.isalnum() or c in (' ', '-', '_')).rstrip().replace(' ', '_') or "roadmap"
        filename = f"roadmap_{safe_domain}_{datetime.now().strftime('%Y%m%d')}.pdf"
        filepath = os.path.join("/tmp", filename)
        pdf.output(filepath)
        return filepath
        
    except Exception as e:
        print(f"PDF Error: {e}")
        return None

# Build Gradio Interface
def create_app():
    with gr.Blocks(title="AI Learning Roadmap Generator") as app:
        roadmap_json = gr.State("")
        domain_name = gr.State("")
        
        gr.Markdown("# 🎯 AI Learning Roadmap Generator")
        gr.Markdown("Generate personalized learning paths for any skill domain")
        
        with gr.Row():
            with gr.Column(scale=1):
                domain_input = gr.Textbox(label="Learning Domain", placeholder="e.g., Machine Learning, Web Development...")
                level_input = gr.Dropdown(["Beginner", "Intermediate", "Advanced"], label="Level", value="Beginner")
                
                with gr.Row():
                    time_input = gr.Number(value=3, minimum=1, label="Time")
                    time_unit = gr.Dropdown(["Weeks", "Months", "Years"], value="Months", label="Unit")
                
                generate_btn = gr.Button("🚀 Generate Roadmap", variant="primary")
                
                gr.Examples(
                    [["Machine Learning", "Beginner", 6, "Months"],
                     ["Web Development", "Intermediate", 4, "Months"],
                     ["Python Programming", "Beginner", 8, "Weeks"]],
                    inputs=[domain_input, level_input, time_input, time_unit]
                )
            
            with gr.Column(scale=2):
                html_output = gr.HTML(label="Your Roadmap")
                
                with gr.Row():
                    download_btn = gr.Button("📄 Download PDF", variant="primary", visible=False)
                
                pdf_output = gr.File(label="Your PDF", interactive=False, visible=False)
        
        def generate(domain, level, time_val, unit):
            if not domain or not domain.strip():
                return "Please enter a domain", gr.update(visible=False), gr.update(visible=False), "", ""
            
            result = generate_roadmap(domain, level, time_val, unit)
            
            if "error" in result:
                return f"<p style='color:red; padding:20px;'>❌ {result['error']}</p>", gr.update(visible=False), gr.update(visible=False), "", ""
            
            html, json_str = format_roadmark_for_display(result)
            if json_str is None:
                return html, gr.update(visible=False), gr.update(visible=False), "", ""
            
            return html, gr.update(visible=True), gr.update(visible=True), json_str, domain
        
        def download(json_data, dom):
            if not json_data:
                return None
            return generate_pdf(json_data, dom)
        
        generate_btn.click(
            fn=generate,
            inputs=[domain_input, level_input, time_input, time_unit],
            outputs=[html_output, download_btn, pdf_output, roadmap_json, domain_name]
        )
        
        download_btn.click(
            fn=download,
            inputs=[roadmap_json, domain_name],
            outputs=pdf_output
        )
    
    return app

app = create_app()
app.launch()