"""Materials Project CLI search tool."""

import os
import sys

from mp_api.client import MPRester
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich.prompt import Prompt, IntPrompt, Confirm
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich import box

from ceraviz.api import render_material
from ceraviz.config import get_stored_api_key, store_api_key
from ceraviz.colors import formula_to_subscript

PAGE_SIZE = 10

console = Console()


def _mask_key(key: str) -> str:
    """Mask an API key for display, showing only first 4 and last 4 chars."""
    if len(key) <= 8:
        return key[:2] + "•" * (len(key) - 2)
    return key[:4] + "•" * (len(key) - 8) + key[-4:]


def setup_api_key() -> str:
    """Interactive API key setup. Returns the resolved key.

    Flow:
    - If a stored key exists: show it masked, offer to use it or change it.
    - If MP_API_KEY env var is set: use it (and offer to store it).
    - Otherwise: prompt for a new key and store it.
    """
    env_key = os.environ.get("MP_API_KEY")
    stored_key = get_stored_api_key()

    console.print()

    # Case 1: Stored key exists
    if stored_key:
        masked = _mask_key(stored_key)
        console.print(f"[green]✓[/green] Stored API key found: [cyan]{masked}[/cyan]")
        change = Confirm.ask("  Change API key?", default=False)

        if not change:
            return stored_key

        # User wants to change
        new_key = Prompt.ask("[cyan]Enter new API key[/cyan]").strip()
        if new_key:
            store_api_key(new_key)
            console.print("[green]✓ API key updated and saved.[/green]")
            return new_key
        else:
            console.print("[yellow]Empty key — keeping existing key.[/yellow]")
            return stored_key

    # Case 2: Env var set but nothing stored
    if env_key:
        masked = _mask_key(env_key)
        console.print(f"[green]✓[/green] Using API key from [cyan]MP_API_KEY[/cyan] env var: [dim]{masked}[/dim]")
        save = Confirm.ask("  Save this key for future sessions?", default=True)
        if save:
            store_api_key(env_key)
            console.print("[green]✓ API key saved.[/green]")
        return env_key

    # Case 3: No key found — first-time setup
    console.print(
        Panel(
            "[yellow]No API key found.[/yellow]\n\n"
            "You need a [bold]free[/bold] Materials Project API key.\n"
            "Get one at: [cyan underline]https://next-gen.materialsproject.org/api[/cyan underline]",
            title="First-Time Setup",
            border_style="yellow",
        )
    )

    new_key = Prompt.ask("\n[cyan]Enter your API key[/cyan]").strip()
    if not new_key:
        console.print("[red]No API key provided. Cannot continue.[/red]")
        sys.exit(1)

    store_api_key(new_key)
    console.print("[green]✓ API key saved for future sessions.[/green]")
    return new_key


def create_results_table(results, start_idx=0, end_idx=None):
    """Create a Rich table displaying material search results."""
    if end_idx is None:
        end_idx = len(results)

    table = Table(
        title=None,
        box=box.ROUNDED,
        header_style="bold cyan",
        border_style="blue",
        row_styles=["", "dim"],
        padding=(0, 1),
    )

    table.add_column("#", style="dim", width=4, justify="right")
    table.add_column("Material ID", style="cyan", width=14)
    table.add_column("Formula", style="bold white", width=20)
    table.add_column("Crystal System", style="green", width=14)
    table.add_column("Band Gap (eV)", style="yellow", width=13, justify="right")

    for i, mat in enumerate(results[start_idx:end_idx], start=start_idx + 1):
        material_id = str(mat.material_id)
        formula_raw = mat.formula_pretty if mat.formula_pretty else "N/A"
        formula = formula_to_subscript(formula_raw) if formula_raw != "N/A" else "N/A"
        crystal_system = mat.symmetry.crystal_system.value if mat.symmetry else "N/A"

        band_gap = mat.band_gap if mat.band_gap is not None else 0.0
        band_gap_str = f"{band_gap:.2f}" if band_gap else "0.00"

        table.add_row(
            str(i),
            material_id,
            formula,
            crystal_system,
            band_gap_str,
        )

    return table


def display_search_menu():
    """Display the main search method selection menu."""
    console.clear()

    title = Text("Materials Project Search", style="bold magenta")
    console.print(Panel(title, border_style="magenta", padding=(1, 2)))
    console.print()

    menu_text = """
[bold cyan]Select a search method:[/bold cyan]

  [yellow]1.[/yellow] [white]Include at least elements[/white]     [dim](e.g., Li,Fe or Si,O,K)[/dim]
  [yellow]2.[/yellow] [white]Include only elements[/white]         [dim](e.g., Li-Fe or Si-O-K)[/dim]
  [yellow]3.[/yellow] [white]Wildcard elements[/white]             [dim](e.g., Li-Fe-*-* or Si-Fe-*-*-*)[/dim]
  [yellow]4.[/yellow] [white]Exact formula[/white]                 [dim](e.g., Li3Fe or Eu2SiCl2O3)[/dim]

  [yellow]q.[/yellow] [red]Quit[/red]
"""
    console.print(Panel(menu_text, border_style="blue", title="Search Options"))


def parse_elements_comma(query):
    """Parse comma-separated elements (include at least)."""
    elements = [e.strip() for e in query.split(',') if e.strip()]
    return elements


def parse_elements_dash(query):
    """Parse dash-separated elements (include only)."""
    elements = [e.strip() for e in query.split('-') if e.strip() and e.strip() != '*']
    return elements


def count_wildcards(query):
    """Count wildcard elements in query."""
    parts = query.split('-')
    return sum(1 for p in parts if p.strip() == '*')


def search_materials(mpr, search_type, query):
    """Execute materials search based on search type."""
    try:
        if search_type == 1:
            elements = parse_elements_comma(query)
            if not elements:
                console.print("[red]No valid elements found in query.[/red]")
                return []

            console.print(f"[dim]Searching for materials containing at least: {', '.join(elements)}[/dim]")
            search_params = {
                "elements": elements,
                "fields": ["material_id", "formula_pretty", "symmetry", "band_gap", "chemsys"]
            }

        elif search_type == 2:
            elements = parse_elements_dash(query)
            if not elements:
                console.print("[red]No valid elements found in query.[/red]")
                return []

            chemsys = "-".join(sorted(elements))
            console.print(f"[dim]Searching for materials with only: {chemsys}[/dim]")
            search_params = {
                "chemsys": chemsys,
                "fields": ["material_id", "formula_pretty", "symmetry", "band_gap", "chemsys"]
            }

        elif search_type == 3:
            elements = parse_elements_dash(query)
            wildcards = count_wildcards(query)

            if not elements:
                console.print("[red]No valid elements found in query.[/red]")
                return []

            total_elements = len(elements) + wildcards
            console.print(f"[dim]Searching for materials with {', '.join(elements)} + {wildcards} wildcard element(s)[/dim]")

            search_params = {
                "elements": elements,
                "num_elements": (total_elements, total_elements) if wildcards > 0 else None,
                "fields": ["material_id", "formula_pretty", "symmetry", "band_gap", "chemsys"]
            }

        elif search_type == 4:
            console.print(f"[dim]Searching for materials with formula: {query}[/dim]")
            search_params = {
                "formula": query,
                "fields": ["material_id", "formula_pretty", "symmetry", "band_gap", "chemsys"]
            }

        else:
            console.print("[red]Invalid search type.[/red]")
            return []

        results = []
        with Progress(
            SpinnerColumn(),
            TextColumn("[cyan]Fetching materials...[/cyan]"),
            BarColumn(),
            TaskProgressColumn(),
            TextColumn("[dim]{task.completed} materials loaded[/dim]"),
            console=console,
            transient=True,
        ) as progress:
            task = progress.add_task("Fetching...", total=None)

            docs_generator = mpr.materials.summary.search(**search_params)

            for doc in docs_generator:
                results.append(doc)
                progress.update(task, completed=len(results))

        console.print(f"[green]✓ Found {len(results)} materials[/green]")
        return results

    except Exception as e:
        console.print(f"[red]Error during search: {e}[/red]")
        return []


def display_results_paginated(results, api_key):
    """Display results with pagination (10 at a time)."""
    if not results:
        console.print(Panel("[yellow]No results found.[/yellow]", border_style="yellow"))
        Prompt.ask("\n[dim]Press Enter to continue[/dim]")
        return

    total = len(results)
    current_page = 0
    total_pages = (total + PAGE_SIZE - 1) // PAGE_SIZE

    while True:
        console.clear()

        start_idx = current_page * PAGE_SIZE
        end_idx = min(start_idx + PAGE_SIZE, total)

        header = Text(f"Search Results", style="bold magenta")
        console.print(Panel(header, border_style="magenta"))
        console.print(f"[dim]Showing {start_idx + 1}-{end_idx} of {total} results (Page {current_page + 1}/{total_pages})[/dim]\n")

        table = create_results_table(results, start_idx, end_idx)
        console.print(table)
        console.print()

        nav_options = []
        if current_page > 0:
            nav_options.append("[yellow]p[/yellow] Previous")
        if current_page < total_pages - 1:
            nav_options.append("[yellow]n[/yellow] Next")
        nav_options.append("[yellow]g[/yellow] Go to page")
        nav_options.append("[yellow]v[/yellow] View crystal structure")
        nav_options.append("[yellow]b[/yellow] Back to menu")

        console.print(" | ".join(nav_options))

        choice = Prompt.ask("\nChoice", default="n" if current_page < total_pages - 1 else "b").lower()

        if choice == 'n' and current_page < total_pages - 1:
            current_page += 1
        elif choice == 'p' and current_page > 0:
            current_page -= 1
        elif choice == 'g':
            try:
                page = IntPrompt.ask(f"Enter page number (1-{total_pages})")
                if 1 <= page <= total_pages:
                    current_page = page - 1
                else:
                    console.print("[red]Invalid page number[/red]")
            except ValueError:
                console.print("[red]Please enter a valid number[/red]")
        elif choice == 'v':
            try:
                row_num = IntPrompt.ask(f"Enter row number to view (1-{total})")
                if 1 <= row_num <= total:
                    material = results[row_num - 1]
                    material_id = str(material.material_id)
                    console.print(f"\n[cyan]Rendering crystal structure for {material_id}...[/cyan]")
                    console.print("[dim]Press 'q' in the PyVista window to close it[/dim]\n")
                    try:
                        render_material(material_id, api_key=api_key)
                    except Exception as e:
                        console.print(f"[red]Error rendering crystal: {e}[/red]")
                        Prompt.ask("[dim]Press Enter to continue[/dim]")
                else:
                    console.print("[red]Invalid row number[/red]")
                    Prompt.ask("[dim]Press Enter to continue[/dim]")
            except ValueError:
                console.print("[red]Please enter a valid number[/red]")
                Prompt.ask("[dim]Press Enter to continue[/dim]")
        elif choice == 'b':
            break


def get_search_prompt(search_type):
    """Get the appropriate prompt based on search type."""
    prompts = {
        1: "Enter elements (comma-separated, e.g., Li,Fe or Si,O,K)",
        2: "Enter elements (dash-separated, e.g., Li-Fe or Si-O-K)",
        3: "Enter elements with wildcards (e.g., Li-Fe-*-* or Si-Fe-*-*-*)",
        4: "Enter exact formula (e.g., Li3Fe or Eu2SiCl2O3)",
    }
    return prompts.get(search_type, "Enter search query")


def main():
    """Main application entry point."""
    # Interactive API key setup (prompts on first run, offers change option)
    api_key = setup_api_key()

    console.print("\n[dim]Connecting to Materials Project...[/dim]")

    try:
        with MPRester(api_key, mute_progress_bars=True) as mpr:
            console.print("[green]✓ Connected successfully![/green]\n")

            while True:
                display_search_menu()

                choice = Prompt.ask("\nEnter your choice").strip().lower()

                if choice == 'q':
                    console.print("\n[magenta]Goodbye![/magenta]")
                    break

                try:
                    search_type = int(choice)
                    if search_type not in [1, 2, 3, 4]:
                        console.print("[red]Please enter 1, 2, 3, 4, or q[/red]")
                        Prompt.ask("[dim]Press Enter to continue[/dim]")
                        continue
                except ValueError:
                    console.print("[red]Please enter 1, 2, 3, 4, or q[/red]")
                    Prompt.ask("[dim]Press Enter to continue[/dim]")
                    continue

                console.print()
                prompt_text = get_search_prompt(search_type)
                query = Prompt.ask(f"[cyan]{prompt_text}[/cyan]")

                if not query.strip():
                    console.print("[red]Empty query. Please try again.[/red]")
                    Prompt.ask("[dim]Press Enter to continue[/dim]")
                    continue

                console.print("\n[dim]Searching...[/dim]")
                results = search_materials(mpr, search_type, query.strip())

                display_results_paginated(results, api_key)

    except Exception as e:
        console.print(f"[red]Failed to connect to Materials Project: {e}[/red]")
        console.print("[yellow]Please check your API key and internet connection.[/yellow]")
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())
