"""
PaperSearchPro - FULL Professional Version
Hierarchical Keyword Search + Export + Responsive UI + BibTeX Support + Author Filter + Specific Paper Search
"""

import requests
import pandas as pd
import time
from datetime import datetime
import gradio as gr  # Make sure this import is present
import os

# ============================================
# CLEAN RESPONSIVE CSS - UPDATED WITH DARKER GRAY BACKGROUND
# ============================================
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

:root {
    --primary: #2563eb;
    --secondary: #7c3aed;
    --border: #d1d5db;
    --bg-soft: #e5e7eb !important;
    --text-main: #111827;
}

/* ================================================= */
/* GENERAL LAYOUT                                   */
/* ================================================= */

body, .gradio-container {
    font-family: 'Inter', sans-serif !important;
    background-color: #e5e7eb !important;
    color: var(--text-main) !important;
}

.gradio-container {
    max-width: 1600px !important;
    margin: auto !important;
}

.white-card {
    background: #ffffff !important;
    border: 1px solid var(--border) !important;
    border-radius: 12px !important;
    padding: 2rem !important;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05) !important;
}

h1 span {
    background: linear-gradient(135deg, var(--primary), var(--secondary));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

.subtitle-text {
    color: #2563eb !important;
    font-weight: 500 !important;
}

/* ================================================= */
/* METRICS                                           */
/* ================================================= */

.metric-value {
    font-size: 2rem !important;
    font-weight: 700 !important;
    color: var(--primary) !important;
}

/* ================================================= */
/* BUTTONS                                           */
/* ================================================= */

button, .gr-button {
    color: #ffffff !important;
    font-weight: 600 !important;
}

button.variant-primary {
    background-color: var(--primary) !important;
}

button.variant-secondary {
    background-color: #6b7280 !important;
}

/* ================================================= */
/* TABS                                              */
/* ================================================= */

.tab-nav {
    background-color: #f3f4f6 !important;
    border-radius: 8px !important;
    padding: 4px !important;
    margin-bottom: 20px !important;
}

.tab-nav button {
    border-radius: 6px !important;
    padding: 8px 16px !important;
    font-weight: 500 !important;
    color: #4b5563 !important;
}

.tab-nav button.selected {
    background-color: var(--primary) !important;
    color: white !important;
}

/* ================================================= */
/* SLIDER COMMENT                                    */
/* ================================================= */

.slider-comment {
    color: #ffffff !important;
    font-size: 0.9rem !important;
    margin-top: 8px !important;
    margin-bottom: 8px !important;
    font-style: italic !important;
    background: #1f2937 !important;
    padding: 8px 12px !important;
    border-radius: 6px !important;
    border-left: 4px solid var(--primary) !important;
    font-weight: 500 !important;
}

/* ================================================= */
/* DROPDOWN BUTTON (Closed State)                   */
/* ================================================= */

.gr-dropdown, .gr-dropdown select, select {
    background-color: #374151 !important;
    color: #ffffff !important;
    border: 1px solid #4b5563 !important;
    border-radius: 6px !important;
}

/* ========================================= */
/* FORCE BUTTON TEXT TO ALWAYS BE WHITE     */
/* ========================================= */

/* All buttons */
button, 
.gr-button {
    color: #ffffff !important;   /* TEXT WHITE */
}

/* Primary button */
button.variant-primary {
    background-color: #2563eb !important;  /* keep blue */
    color: #ffffff !important;             /* TEXT WHITE */
}

/* Secondary button */
button.variant-secondary {
    background-color: #6b7280 !important;  /* keep gray */
    color: #ffffff !important;             /* TEXT WHITE */
}

/* Keep text white even on hover */
button:hover {
    color: #ffffff !important;
}
/* ================================================= */
/* METRIC LABEL BLACK BACKGROUND                    */
/* ================================================= */

.metric-label {
    display: inline-block !important;
    background-color: #111827 !important;  /* Black */
    color: #ffffff !important;             /* White text */
    padding: 6px 12px !important;
    border-radius: 6px !important;
    font-weight: 600 !important;
}



/* ================================================= */
/* METRIC VALUE BLACK BACKGROUND FIX                */
/* ================================================= */

.metric-value {
    display: inline-block !important;
    background-color: #111827 !important;  /* Black */
    color: #ffffff !important;             /* White text */
    padding: 6px 12px !important;
    border-radius: 6px !important;
    font-weight: 700 !important;
}

/* Remove default white background from strong */
.metric-value strong {
    background: transparent !important;
    color: #ffffff !important;
}
"""
# ============================================
# SEARCH ENGINE
# ============================================

class PaperSearchPro:

    def __init__(self):
        self.base_url = "https://api.openalex.org/works"
        self.per_page = 200

    def search_papers(self, query, keywords, authors, max_papers, year_from, year_to, open_access, sort_by, sort_order):

        filters = []
        if year_from:
            filters.append(f"publication_year:>{int(year_from)-1}")
        if year_to:
            filters.append(f"publication_year:<{int(year_to)+1}")
        if open_access:
            filters.append("open_access.is_oa:true")

        # Add author filter if provided
        if authors and authors.strip():
            # Clean and format author names for filtering
            author_list = [a.strip() for a in authors.split(",") if a.strip()]
            if author_list:
                # OpenAlex uses author.search filter
                author_filter = "|".join(author_list)
                filters.append(f"author.search:{author_filter}")

        filter_str = ",".join(filters) if filters else None

        # Map sort options to API parameters
        sort_field_map = {
            "Relevance": "relevance_score",
            "Year": "publication_date",
            "Citations": "cited_by_count"
        }

        sort_field = sort_field_map.get(sort_by, "relevance_score")
        sort_direction = "desc" if sort_order == "Descending" else "asc"
        sort_param = f"{sort_field}:{sort_direction}"

        all_papers = []
        page = 1

        while len(all_papers) < max_papers:
            params = {
                "search": query,
                "per-page": self.per_page,
                "page": page,
                "sort": sort_param
            }

            if filter_str:
                params["filter"] = filter_str

            try:
                response = requests.get(self.base_url, params=params, timeout=20)
                data = response.json()
                results = data.get("results", [])

                if not results:
                    break

                for work in results:
                    paper = self.extract_info(work)
                    all_papers.append(paper)

                page += 1
                time.sleep(0.2)
            except Exception as e:
                print(f"Error: {e}")
                break

        # ===============================
        # HIERARCHICAL KEYWORD FILTERING
        # ===============================

        keyword_list = [k.strip().lower() for k in keywords.split(",") if k.strip()]

        if len(keyword_list) >= 2:
            filtered = []

            for p in all_papers:
                combined = f"{p['Title']} {p['Abstract']} {p['Authors']} {p['All Authors']}".lower()

                # first 2 keywords REQUIRED
                if keyword_list[0] not in combined or keyword_list[1] not in combined:
                    continue

                score = 2
                matched = [keyword_list[0], keyword_list[1]]

                for kw in keyword_list[2:]:
                    if kw in combined:
                        score += 1
                        matched.append(kw)

                p["Relevance"] = score
                p["Matched Keywords"] = ", ".join(matched)

                filtered.append(p)

            # Sort by relevance score (our calculated score)
            filtered.sort(key=lambda x: x["Relevance"], reverse=True)
            all_papers = filtered[:max_papers]

        # ===============================
        # AUTHOR FILTERING (post-API)
        # ===============================
        elif authors and authors.strip():
            # If we have author filter but no keyword hierarchy
            author_list = [a.strip().lower() for a in authors.split(",") if a.strip()]
            filtered = []

            for p in all_papers:
                all_authors_lower = p['All Authors'].lower()
                match_all = True
                matched_authors = []

                for author in author_list:
                    if author in all_authors_lower:
                        matched_authors.append(author)
                    else:
                        match_all = False
                        break

                if match_all:
                    p["Relevance"] = len(matched_authors)
                    p["Matched Keywords"] = f"Authors: {', '.join(matched_authors)}"
                    filtered.append(p)

            all_papers = filtered[:max_papers]

        total = len(all_papers)
        avg_rel = sum(p.get("Relevance", 0) for p in all_papers) / total if total else 0
        years = [p["Year"] for p in all_papers if p["Year"] != "N/A"]
        year_range = f"{min(years)} - {max(years)}" if years else "N/A"

        df = pd.DataFrame(all_papers)

        # Reorder columns to include BibTeX and remove All Authors from display
        column_order = ['Title', 'Authors', 'Year', 'Citations', 'Journal', 'Abstract',
                       'BibTeX', 'Relevance', 'Matched Keywords']
        df = df[column_order]

        return df, all_papers, total, avg_rel, year_range

    def search_specific_paper(self, search_type, search_value):
        """Search for a specific paper by title, author, or DOI"""

        all_papers = []

        if search_type == "DOI":
            # Clean DOI (remove URL if present)
            if "doi.org/" in search_value:
                search_value = search_value.split("doi.org/")[-1]

            # Direct lookup by DOI
            url = f"https://api.openalex.org/works/doi/{search_value}"
            try:
                response = requests.get(url, timeout=15)
                if response.status_code == 200:
                    work = response.json()
                    paper = self.extract_info(work)
                    all_papers.append(paper)
            except Exception as e:
                print(f"Error fetching DOI: {e}")

        elif search_type == "Exact Title":
            # Search by exact title
            params = {
                "search": f'"{search_value}"',  # Quote for exact match
                "per-page": 10,
                "sort": "relevance_score:desc"
            }
            try:
                response = requests.get(self.base_url, params=params, timeout=15)
                data = response.json()
                results = data.get("results", [])

                for work in results:
                    # Check for close title match
                    title = work.get("title", "").lower()
                    if search_value.lower() in title or title in search_value.lower():
                        paper = self.extract_info(work)
                        all_papers.append(paper)
                        break  # Take the first good match
            except Exception as e:
                print(f"Error searching title: {e}")

        elif search_type == "Specific Author":
            # Search for papers by a specific author
            params = {
                "search": search_value,
                "filter": "authorships.author.display_name.search:" + search_value,
                "per-page": 50,
                "sort": "cited_by_count:desc"
            }
            try:
                response = requests.get(self.base_url, params=params, timeout=15)
                data = response.json()
                results = data.get("results", [])

                for work in results[:20]:  # Limit to 20 papers
                    paper = self.extract_info(work)
                    all_papers.append(paper)
            except Exception as e:
                print(f"Error searching author: {e}")

        if all_papers:
            for i, p in enumerate(all_papers):
                p['Rank'] = i + 1

            df = pd.DataFrame(all_papers)
            column_order = ['Title', 'Authors', 'Year', 'Citations', 'Journal', 'Abstract', 'BibTeX']
            # Only include columns that exist
            available_cols = [col for col in column_order if col in df.columns]
            df = df[available_cols]

            summary = f"✅ Found **{len(all_papers)}** paper(s) matching your query"
            return df, all_papers, summary

        return pd.DataFrame(), [], "❌ No papers found"

    def extract_info(self, work):
        # Rebuild abstract
        abstract = ""
        if work.get("abstract_inverted_index"):
            words = []
            for word, positions in work["abstract_inverted_index"].items():
                for pos in positions:
                    words.append((pos, word))
            words.sort()
            abstract = " ".join([w[1] for w in words])
        else:
            abstract = "No abstract available"

        # Get all authors - FULL LIST (no et al.)
        all_authors = []

        for authorship in work.get("authorships", []):
            # FIRST PRIORITY: Use raw_author_name (full name as published)
            raw_name = authorship.get("raw_author_name", "")
            display_name = authorship.get("author", {}).get("display_name", "")

            # Choose the best name
            if raw_name and len(raw_name) > 2:
                # Raw name is usually more complete (e.g., "Mamoru SHOJI" vs "M. Shoji")
                author_name = raw_name
            elif display_name:
                author_name = display_name
            else:
                author_name = ""

            if author_name:
                all_authors.append(author_name)

        # Show ALL authors, no truncation
        authors_str = ", ".join(all_authors)

        # Get DOI for reference
        doi = work.get("doi", "")
        if doi:
            doi = doi.replace("https://doi.org/", "")

        # Generate BibTeX entry
        bibtex = self.generate_bibtex(work, all_authors)

        return {
            "Title": work.get("title", "No title"),
            "Authors": authors_str,
            "All Authors": ", ".join(all_authors),  # Hidden field for filtering
            "Year": work.get("publication_year", "N/A"),
            "Citations": work.get("cited_by_count", 0),
            "Journal": (work.get("primary_location", {}).get("source") or {}).get("display_name", "Unknown"),
            "DOI": doi,
            "Abstract": abstract,
            "BibTeX": bibtex,
            "Relevance": 0,
            "Matched Keywords": ""
        }

    def generate_bibtex(self, work, authors):
        """Generate BibTeX entry for a paper"""

        # Create citation key
        first_author = authors[0].split()[-1] if authors else "Unknown"
        year = work.get("publication_year", "0000")
        title_words = work.get("title", "Untitled").split()[:3]
        title_short = "".join([w[0] for w in title_words if w]) if title_words else "Paper"
        citation_key = f"{first_author}{year}{title_short}"

        # Clean the citation key (remove special characters)
        citation_key = ''.join(c for c in citation_key if c.isalnum())

        # Get DOI or URL
        doi = work.get("doi", "")
        url = work.get("primary_location", {}).get("landing_page_url", "")
        if not url and work.get("open_access", {}).get("oa_url"):
            url = work.get("open_access", {}).get("oa_url")

        # Determine entry type (article, book, etc.)
        pub_type = "article"
        if work.get("type") == "book":
            pub_type = "book"
        elif work.get("type") == "book-chapter":
            pub_type = "incollection"

        # Build BibTeX string
        bibtex_lines = [
            f"@{pub_type}{{{citation_key},",
            f"  title = {{{work.get('title', 'Untitled')}}},",
            f"  author = {{{' and '.join(authors)}}},",
            f"  year = {{{work.get('publication_year', '0000')}}},"
        ]

        # Add journal/booktitle
        source = (work.get("primary_location", {}).get("source") or {})
        if pub_type == "article":
            journal = source.get("display_name", "Unknown Journal")
            bibtex_lines.append(f"  journal = {{{journal}}},")
        else:
            booktitle = source.get("display_name", "Unknown Publication")
            bibtex_lines.append(f"  booktitle = {{{booktitle}}},")

        # Add volume, issue if available
        if work.get("biblio", {}).get("volume"):
            bibtex_lines.append(f"  volume = {{{work['biblio']['volume']}}},")
        if work.get("biblio", {}).get("issue"):
            bibtex_lines.append(f"  number = {{{work['biblio']['issue']}}},")

        # Add pages if available
        if work.get("biblio", {}).get("first_page") and work.get("biblio", {}).get("last_page"):
            pages = f"{work['biblio']['first_page']}--{work['biblio']['last_page']}"
            bibtex_lines.append(f"  pages = {{{pages}}},")

        # Add DOI/URL
        if doi:
            bibtex_lines.append(f"  doi = {{{doi}}},")
        elif url:
            bibtex_lines.append(f"  url = {{{url}}},")

        # Add abstract if available
        if work.get("abstract_inverted_index"):
            # Rebuild abstract for BibTeX
            words = []
            for word, positions in work["abstract_inverted_index"].items():
                for pos in positions:
                    words.append((pos, word))
            words.sort()
            abstract = " ".join([w[1] for w in words])
            if abstract:
                bibtex_lines.append(f"  abstract = {{{abstract}}},")

        # Close the entry
        bibtex_lines.append("}")

        return "\n".join(bibtex_lines)


# ============================================
# INTERFACE
# ============================================

def create_interface():

    with gr.Blocks(
    title="PaperSearchPro",
    css=custom_css,
    theme=gr.themes.Default(
        primary_hue="blue",
        neutral_hue="gray"
    )) as demo:

        gr.HTML("""
        <div style='text-align:center; padding:3rem 0'>
            <h1><span>📱XYZ:📱PaperSearchPro</span></h1>
            <p class='subtitle-text'>Easy Academic Paper Searching tool by Paul</p>
        </div>
        """)

        # Create tabs for different search modes
        with gr.Tabs(elem_classes="tab-nav"):
            # Tab 1: General Search
            with gr.TabItem("🔍 General Search", id=0):
                with gr.Column(elem_classes="white-card"):

                    query = gr.Textbox(
                        label="Search Topic",
                        value="photoacoustic imaging",
                        placeholder="Enter research topic or keywords..."
                    )

                    with gr.Row():
                        with gr.Column():
                            keywords = gr.Textbox(
                                label="Required Keywords (first 2 REQUIRED for hierarchical search)",
                                value="blood, photoacoustics, imaging",
                                placeholder="e.g., blood, photoacoustics, imaging (comma-separated)"
                            )

                        with gr.Column():
                            authors = gr.Textbox(
                                label="Filter by Authors (optional)",
                                value="",
                                placeholder="e.g., Smith, Zhang, Wang (comma-separated)"
                            )

                    # ==========================
                    # MAX PAPERS SLIDER WITH BLACK REGION COMMENT
                    # ==========================
                    max_papers = gr.Slider(1000, 50000, value=1000, step=100, label="Max Papers")

                    # Add black region comment below slider
                    gr.HTML("""
                    <div class="slider-comment">
                        ⚡ Tip: Set this number as high as possible or more than 5000 for better scanning of all available papers
                    </div>
                    """)

                    with gr.Row():
                        year_from = gr.Number(value=2000, label="Year From", precision=0)
                        year_to = gr.Number(value=2026, label="Year To", precision=0)
                        open_access = gr.Checkbox(label="Open Access Only")

                    # ==========================
                    # SORTING OPTIONS
                    # ==========================
                    with gr.Row():
                        sort_by = gr.Dropdown(
                            label="Sort By",
                            choices=["Relevance", "Year", "Citations"],
                            value="Relevance"
                        )
                        sort_order = gr.Dropdown(
                            label="Sort Order",
                            choices=["Descending", "Ascending"],
                            value="Descending"
                        )

                    search_btn = gr.Button("🔍 SEARCH PAPERS", variant="primary", size="lg")

            # Tab 2: Specific Paper Search
            with gr.TabItem("🎯 Find Specific Paper", id=1):
                with gr.Column(elem_classes="white-card"):
                    gr.Markdown("""
                    ### Find a specific paper by:
                    - **DOI** (e.g., 10.1038/nature12373)
                    - **Exact Title** (e.g., "Photoacoustic imaging in biomedicine")
                    - **Author** (find papers by a specific researcher)
                    """)

                    with gr.Row():
                        specific_search_type = gr.Dropdown(
                            label="Search by",
                            choices=["DOI", "Exact Title", "Specific Author"],
                            value="DOI"
                        )

                        specific_search_value = gr.Textbox(
                            label="Enter value",
                            placeholder="e.g., 10.1038/nature12373 or exact paper title or author name",
                            scale=2
                        )

                    specific_search_btn = gr.Button("🔍 FIND PAPER", variant="primary")
                    specific_summary = gr.Markdown("Ready")

            # Metrics section (common for both tabs)
            with gr.Row():
                total_display = gr.Markdown("**0**", elem_classes="metric-value")
                avg_display = gr.Markdown("**0.0**", elem_classes="metric-value")
                year_display = gr.Markdown("**N/A**", elem_classes="metric-value")

            # Add labels for metrics
            with gr.Row():
                gr.Markdown("**Total Papers**", elem_classes="metric-label")
                gr.Markdown("**Avg Relevance**", elem_classes="metric-label")
                gr.Markdown("**Year Range**", elem_classes="metric-label")

            # Results section (common for both tabs)
            results = gr.Dataframe(
                label="Results",
                row_count=15,
                column_widths=["25%", "15%", "5%", "5%", "10%", "20%", "20%", "5%", "5%"],
                wrap=True
            )
            state = gr.State([])

            # ==========================
            # EXPORT SECTION
            # ==========================

            gr.HTML("""
            <h3 style="color: #2563eb !important; font-weight: 600; margin-bottom: 16px; margin-top: 24px;">📥 Export Results</h3>
            """)

            with gr.Row():
                filename = gr.Textbox(
                    label="Filename",
                    value="papers_results",
                    placeholder="Enter base filename without extension"
                )
                format_type = gr.Dropdown(
                    label="Format",
                    choices=["TXT", "CSV", "Excel", "BibTeX"],
                    value="TXT"
                )

            export_btn = gr.Button("📥 Download Results", variant="secondary")

            export_status = gr.Markdown("")
            download_file = gr.File(visible=False)

        searcher = PaperSearchPro()

        def search_wrapper(q, k, a, m, yf, yt, oa, sb, so):
            try:
                df, papers, total, avg, yr = searcher.search_papers(q, k, a, m, yf, yt, oa, sb, so)
                return df, papers, f"**{total}**", f"**{avg:.2f}**", f"**{yr}**"
            except Exception as e:
                # Return empty results with error message in status
                empty_df = pd.DataFrame(columns=['Title', 'Authors', 'Year', 'Citations', 'Journal', 'Abstract', 'BibTeX', 'Relevance', 'Matched Keywords'])
                return empty_df, [], "**0**", "**0.0**", "**N/A**"

        def specific_search_wrapper(search_type, search_value):
            try:
                df, papers, summary = searcher.search_specific_paper(search_type, search_value)
                if not papers:
                    return df, papers, summary, "**0**", "**0.0**", "**N/A**"

                # Calculate metrics for display
                total = len(papers)
                avg_rel = 0
                years = [p["Year"] for p in papers if p["Year"] != "N/A"]
                year_range = f"{min(years)} - {max(years)}" if years else "N/A"

                return df, papers, summary, f"**{total}**", f"**{avg_rel:.2f}**", f"**{year_range}**"
            except Exception as e:
                empty_df = pd.DataFrame(columns=['Title', 'Authors', 'Year', 'Citations', 'Journal', 'Abstract', 'BibTeX'])
                return empty_df, [], f"❌ Error: {str(e)}", "**0**", "**0.0**", "**N/A**"

        # Connect search buttons
        search_btn.click(
            search_wrapper,
            inputs=[query, keywords, authors, max_papers, year_from, year_to, open_access, sort_by, sort_order],
            outputs=[results, state, total_display, avg_display, year_display]
        )

        specific_search_btn.click(
            specific_search_wrapper,
            inputs=[specific_search_type, specific_search_value],
            outputs=[results, state, specific_summary, total_display, avg_display, year_display]
        )

        def export_results(papers, fname, fmt):

            if not papers:
                return None, "❌ No papers to export", gr.File(visible=False)

            df = pd.DataFrame(papers)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

            # Remove All Authors column from export if it exists
            if 'All Authors' in df.columns:
                df = df.drop(columns=['All Authors'])

            if fmt == "CSV":
                path = f"{fname}_{timestamp}.csv"
                df.to_csv(path, index=False)

            elif fmt == "Excel":
                path = f"{fname}_{timestamp}.xlsx"
                df.to_excel(path, index=False, engine='openpyxl')

            elif fmt == "BibTeX":
                path = f"{fname}_{timestamp}.bib"
                with open(path, "w", encoding="utf-8") as f:
                    for _, row in df.iterrows():
                        if row['BibTeX'] and row['BibTeX'] != "":
                            f.write(row['BibTeX'] + "\n\n")

            else:  # TXT format
                path = f"{fname}_{timestamp}.txt"
                with open(path, "w", encoding="utf-8") as f:
                    f.write("PAPER SEARCH PRO RESULTS\n")
                    f.write("="*80 + "\n")
                    f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                    f.write(f"Total Papers: {len(df)}\n")
                    f.write("="*80 + "\n\n")

                    for i, row in df.iterrows():
                        f.write(f"PAPER #{i+1}\n")
                        f.write("-"*40 + "\n")
                        for col in df.columns:
                            if col == "BibTeX":
                                f.write(f"\n{col}:\n{row[col]}\n")
                            else:
                                f.write(f"{col}: {row[col]}\n")
                        f.write("\n" + "="*80 + "\n\n")

            return path, f"✅ File ready: {os.path.basename(path)}", gr.File(visible=True)

        export_btn.click(
            export_results,
            inputs=[state, filename, format_type],
            outputs=[download_file, export_status, download_file]
        )

    return demo


# ============================================
# MAIN
# ============================================

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        inbrowser=True,
        height=1000,
        server_name="0.0.0.0",  # Makes it accessible on network
        share=False  # Set to True for temporary public link
    )