import gradio as gr
from pathlib import Path
import base64

# Basic paths
ROOT = Path(__file__).resolve().parent

# Hardcoded logo as base64 (cannot be downloaded)
LOGO_BASE64 = None
try:
    with open(ROOT / "paper2agent_logo.txt", "rb") as f:
        LOGO_BASE64 = f.read().decode()
except:
    pass

# =====================
# Gradio UI Layout Only
# =====================
with gr.Blocks(title="Paper2Agent") as iface:
    # Logo at top left (hardcoded, cannot be downloaded)
    if LOGO_BASE64:
        gr.HTML(f'<img src="data:image/png;base64,{LOGO_BASE64}" style="height:80px;width:auto;" />')

    gr.Markdown("""
[Paper](https://arxiv.org/abs/2509.06917) | [GitHub](https://github.com/jmiao24/Paper2Agent)

**TL;DR:** Upload your paper code repo and get an auto-generated mcp.
Please be patient — takes about 20–30 minutes to process.
""", elem_id="intro-md")

    # -------- Input/Output Layout --------
    with gr.Row():
        # ========== LEFT: INPUT ==========
        with gr.Column(scale=1):
            with gr.Accordion("Input", open=True):
                github_in = gr.Textbox(
                    label="📘 GitHub Repo URL",
                    placeholder="https://github.com/google-deepmind/alphagenome"
                )

                key_in = gr.Textbox(
                    label="🔑 Claude API Key",
                    placeholder="sk-ant-...",
                    type="password"
                )

                repo_key_in = gr.Textbox(
                    label="🔐 API Key (optional, for repositories requiring authentication)",
                    placeholder="Enter API key for private repositories",
                    type="password"
                )

                tutorials_in = gr.Textbox(
                    label="📚 Tutorials (optional)",
                    placeholder="Filter tutorials by title or URL"
                )

                run_btn = gr.Button("🚀 Run", variant="primary")
                example_btn = gr.Button("📝 Use Example Values", variant="secondary")

                # Example values info
                gr.Markdown("""
                <details style="margin-top: 8px; padding: 10px; background: #f8f9fa; border-radius: 6px; border: 1px solid #e0e0e0;">
                <summary style="cursor: pointer; font-weight: bold; color: #333;">💡 Example Values</summary>
                <div style="margin-top: 10px; font-size: 0.85em;">
                <div style="margin-bottom: 8px;">
                    <div style="font-weight: bold; margin-bottom: 2px;">GitHub URL:</div>
                    <code style="display: block; background: white; padding: 6px 8px; border-radius: 4px; border: 1px solid #ddd; word-break: break-all;">https://github.com/google-deepmind/alphagenome</code>
                </div>
                <div style="margin-bottom: 8px;">
                    <div style="font-weight: bold; margin-bottom: 2px;">Claude API Key:</div>
                    <code style="display: block; background: white; padding: 6px 8px; border-radius: 4px; border: 1px solid #ddd; word-break: break-all;">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</code>
                </div>
                <div>
                    <div style="font-weight: bold; margin-bottom: 2px;">Repo API Key:</div>
                    <code style="display: block; background: white; padding: 6px 8px; border-radius: 4px; border: 1px solid #ddd; word-break: break-all;">AIzaSyDZ-IxStzMSUElDGWS7U9v6BIDr_0WMoO8</code>
                </div>
                </div>
                </details>
                """, elem_id="example-section")

        # ========== RIGHT: OUTPUT ==========
        with gr.Column(scale=1):
            with gr.Accordion("Output", open=True):
                # Logs with scrolling enabled
                logs_out = gr.Textbox(
                    label="🧾 Logs",
                    lines=20,
                    max_lines=20,
                    autoscroll=False
                )
                # Downloads
                with gr.Row():
                    zip_out = gr.File(
                        label="📦 Download Results (.zip)",
                        interactive=False,
                        visible=True,
                        scale=1
                    )
                overleaf_out = gr.HTML(label="Open in Overleaf")

    # Fill example values
    def fill_example():
        return (
            "https://github.com/google-deepmind/alphagenome",
            "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",
            "AIzaSyDZ-IxStzMSUElDGWS7U9v6BIDr_0WMoO8",
            ""
        )

    # Button click handler
    def run_pipeline(github_url, repo_api_key, claude_api_key, tutorials_filter):
        """
        Run the Paper2Agent pipeline with the provided inputs.
        """
        import subprocess
        import os

        ui_logs = []  # Simplified logs for UI

        try:
            # Validate inputs
            if not github_url or not github_url.strip():
                ui_logs.append("❌ Error: GitHub Repo URL is required")
                return "\n".join(ui_logs), None, ""

            if not claude_api_key or not claude_api_key.strip():
                ui_logs.append("❌ Error: Claude API Key is required")
                return "\n".join(ui_logs), None, ""

            # Create Results folder
            results_path = ROOT / "Results"
            results_path.mkdir(parents=True, exist_ok=True)

            ui_logs.append(f"🚀 Starting Paper2Agent pipeline...")
            ui_logs.append(f"📘 GitHub Repo: {github_url}")
            ui_logs.append(f"🔑 Claude API Key: {'*' * (len(claude_api_key) - 4)}{claude_api_key[-4:]}")

            if tutorials_filter:
                ui_logs.append(f"📚 Tutorial Filter: {tutorials_filter}")

            ui_logs.append(f"\n📝 Detailed logs will be saved to: Results/log.log")
            ui_logs.append("\n" + "="*70)
            yield "\n".join(ui_logs), None, ""

            # Set environment variable for Claude API key (for SDK initialization)
            env = os.environ.copy()
            env['ANTHROPIC_API_KEY'] = claude_api_key
            env['PYTHONUNBUFFERED'] = '1'

            # Build command with unbuffered Python
            cmd = [
                "python", "-u", "test.py",
                "--github_url", github_url
            ]

            # Add repo API key if provided (for repository authentication)
            if repo_api_key and repo_api_key.strip():
                cmd.extend(["--api", repo_api_key])

            if tutorials_filter and tutorials_filter.strip():
                cmd.extend(["--tutorials", tutorials_filter])

            # Run test.py and capture stdout for UI
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=0,
                env=env
            )

            # Stream output to UI
            for line in iter(process.stdout.readline, ''):
                if line:
                    stripped_line = line.strip()
                    if stripped_line:
                        ui_logs.append(stripped_line)
                        yield "\n".join(ui_logs), None, ""

            process.wait()

            if process.returncode == 0:
                ui_logs.append("\n" + "="*70)
                ui_logs.append("✅ Pipeline completed successfully!")
                ui_logs.append("="*70)

                # Create zip file from Results folder
                zip_file = None

                ui_logs.append("\n📦 Creating zip archive from Results folder...")

                if results_path.exists():
                    import shutil

                    # Create zip file with timestamp
                    zip_base_name = f"Results"
                    zip_file_path = ROOT / zip_base_name

                    try:
                        # Create zip archive of the entire Results folder
                        shutil.make_archive(
                            str(zip_file_path),
                            'zip',
                            ROOT,
                            'Results'
                        )
                        zip_file = str(zip_file_path) + ".zip"
                        ui_logs.append(f"✅ Created zip file: {zip_file}")
                        ui_logs.append(f"📥 Ready for download!")
                        ui_logs.append(f"\n📝 Full logs saved to: Results/log.log")
                        yield "\n".join(ui_logs), zip_file, ""
                    except Exception as e:
                        ui_logs.append(f"⚠️ Failed to create zip file: {str(e)}")
                        yield "\n".join(ui_logs), None, ""
                else:
                    ui_logs.append(f"⚠️ Results folder not found at: {results_path}")
                    yield "\n".join(ui_logs), None, ""
            else:
                ui_logs.append("\n" + "="*70)
                ui_logs.append(f"❌ Pipeline failed with exit code {process.returncode}")
                ui_logs.append(f"📝 Check logs for details: Results/log.log")
                ui_logs.append("="*70)
                yield "\n".join(ui_logs), None, ""

        except Exception as e:
            ui_logs.append(f"\n❌ Error: {str(e)}")
            ui_logs.append(f"📝 Check logs for details: Results/log.log")
            yield "\n".join(ui_logs), None, ""

    # Connect example button
    example_btn.click(
        fn=fill_example,
        inputs=[],
        outputs=[github_in, key_in, repo_key_in, tutorials_in]
    )

    # Connect run button
    run_btn.click(
        fn=run_pipeline,
        inputs=[github_in, repo_key_in, key_in, tutorials_in],
        outputs=[logs_out, zip_out, overleaf_out]
    )

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