import typer
import sys
import time
import os
from dotenv import load_dotenv
from rich.prompt import Prompt
from rich.console import Console
from rich import print

# Load .env file if present
load_dotenv()
from .auth import login as auth_login, logout as auth_logout, is_authenticated, get_config
from .api import fetch_user_credits, generate_schema, check_schema_status, create_database, save_db_credentials, delete_database, fetch_user_projects, save_user_project, delete_user_project, append_project_query, deduct_user_credit
from .schema import build_schema, review_schema_options

app = typer.Typer(help="AIGenDB CLI Platform")
console = Console()

def check_credits_and_start():
    with console.status("Checking credits and projects..."):
        ui_data = fetch_user_credits()
        projects = fetch_user_projects()
        
    print(f"Remaining Credits: [yellow]{ui_data.get('credits', 0)}[/yellow]")
    print(f"Account Type: [blue]{ui_data.get('subscription_type', 'Free')}[/blue]")
    project_names = []
    if projects:
        project_names = [p.get("project_name", "Unknown") for p in projects]
        print(f"Existing Projects: [cyan]{', '.join(project_names)}[/cyan]\n")
    else:
        print("Existing Projects: [cyan]None[/cyan]\n")
    
    if ui_data.get('credits', 0) <= 0:
        print("[red]You have no credits remaining.[/red]")
        print("Please upgrade your account or Request more credits to continue.")
        print("[blue]https://www.aigendb.xyz/pricing[/blue]\n")
        raise typer.Exit()
        
    schema_payload = build_schema(project_names)
    
    with console.status("Generating schema from ML server... (This may take a moment)"):
        generate_res = generate_schema(schema_payload)
        project_name = generate_res.get("project", schema_payload.get("project_name"))
        sql_queries = generate_res.get("sql_queries", [])
    sql_result = "\n".join(sql_queries)
    
    if not sql_result.strip():
        print("[red]\n✗ Schema generation failed — the ML server did not respond.[/red]")
        print("[yellow]The server may be busy or starting up. Please wait a moment and try again.\n[/yellow]")
        raise typer.Exit()
    
    print("[green]SQL generation completed\n[/green]")
    
    with console.status("Saving project to dashboard..."):
        save_user_project(project_name, schema_payload, sql_result)
    print("")
    
    print("[cyan]--- Generated Schema ---[/cyan]")
    if sql_result:
        print(sql_result)
    else:
        print("[italic]Schema generation placeholder string returned from backend.[/italic]")

    print("[cyan]------------------------[/cyan]\n")
    
    choice = review_schema_options()
    
    if choice == "cancel":
        with console.status("Deleting project..."):
            delete_user_project(project_name)
        print("[red]Schema generation cancelled. Project deleted.[/red]\n")
        raise typer.Exit()
        
    if choice == "modify":
        print("\n[cyan]Opening schema editor...[/cyan]")
        print("[italic]Make your changes, save the file, and close the editor to continue.[/italic]")
        
        new_sql = typer.edit(text=sql_result)
        final_sql = sql_result  # default to original if no changes
        
        if new_sql and new_sql.strip() != sql_result.strip():
            with console.status("Saving modified schema..."):
                success = append_project_query(project_name, new_sql)
            if success:
                print("[green]Schema modification saved.[/green]")
                final_sql = new_sql
            else:
                print("[red]Failed to save modified schema. Proceeding with original.[/red]")
        else:
            print("[yellow]No modifications detected. Using original schema.[/yellow]")

        # Provision database with the final SQL
        print("\n[cyan]Provisioning your database...[/cyan]")
        creds = None
        while not creds:
            with console.status("Connecting to database provisioner... (This may take a moment)"):
                creds = create_database(project_name, final_sql)
            if not creds:
                print("[red]\n✗ Database provisioning failed.[/red]")
                retry = Prompt.ask("What would you like to do?", choices=["retry", "cancel"],
                                   default="retry",
                                   show_choices=True,
                                   show_default=True)
                if retry == "cancel":
                    print("[yellow]Provisioning cancelled. Your schema change is saved safely.\n[/yellow]")
                    raise typer.Exit()
                print("[cyan]Retrying...\n[/cyan]")

        db_name = creds.get("database_name", "")
        conn_url = creds.get("connection_url", "")
        db_pass  = creds.get("password", "")

        with console.status("Saving database credentials..."):
            save_db_credentials(project_name, db_name, conn_url, db_pass)
            deduct_user_credit(project_name)

        p_list = fetch_user_projects()
        p_obj = next((p for p in p_list if p.get("project_name") == project_name), {})
        project_id = p_obj.get("project_id", project_name)
        url = f"https://www.aigendb.xyz/dashboard/projects/{project_id}"
        print(f"\n[green]Database is Created successfully, Acess the Database Credentials in you account dashboard : [/green]")
        print(f"[blue][u]{url}[/u][/blue]\n")
        typer.launch(url)
        raise typer.Exit()
        
    if choice == "accept":
        print("\n[cyan]Provisioning your database...[/cyan]")
        creds = None
        while not creds:
            with console.status("Connecting to database provisioner... (This may take a moment)"):
                creds = create_database(project_name, sql_result)
            if not creds:
                print("[red]\n✗ Database provisioning failed.[/red]")
                retry = Prompt.ask("What would you like to do?", choices=["retry", "cancel"],
                                   default="retry",
                                   show_choices=True,
                                   show_default=True)
                if retry == "cancel":
                    print("[yellow]Provisioning cancelled. Your schema is saved safely.\n[/yellow]")
                    raise typer.Exit()
                print("[cyan]Retrying...\n[/cyan]")

        db_name = creds.get("database_name", "")
        conn_url = creds.get("connection_url", "")
        db_pass  = creds.get("password", "")

        with console.status("Saving database credentials..."):
            save_db_credentials(project_name, db_name, conn_url, db_pass)
            deduct_user_credit(project_name)

        p_list = fetch_user_projects()
        p_obj = next((p for p in p_list if p.get("project_name") == project_name), {})
        project_id = p_obj.get("project_id", project_name)
        url = f"https://www.aigendb.xyz/dashboard/projects/{project_id}"
        print(f"\n[green]Database is Created successfully, Acess the Database Credentials in you account dashboard : [/green]")
        print(f"[blue][u]{url}[/u][/blue]\n")
        typer.launch(url)
        raise typer.Exit()


def cleanup_orphaned_projects(projects):
    """Silently delete any projects that have no valid schema (failed/aborted generations)."""
    for project in projects:
        queries = project.get("queries", {})
        has_valid_schema = any(
            v.strip() for v in queries.values()
        ) if queries else False
        if not has_valid_schema:
            delete_user_project(project.get("project_name", ""))


@app.callback(invoke_without_command=True)
def main(ctx: typer.Context):
    if ctx.invoked_subcommand is None:
        if is_authenticated():
            config = get_config()
            banner = """
[cyan]    _   ___  ___          ___  ___ [/cyan]
[cyan]   /_\\ |_ _|/ __| ___ _ _|   \\| _ )[/cyan]
[cyan]  / _ \\ | || (_ |/ -_) ' \\ |) | _ \\[/cyan]
[cyan] /_/ \\_\\___|\\___|\\___|_||_|___/___/[/cyan]
"""
            print(banner)
            print("[cyan]Welcome back![/cyan]")
            if config and config.get("email"):
                print(f"Logged in as: [green]{config.get('email')}[/green]")
            
            projects = fetch_user_projects()
            cleanup_orphaned_projects(projects)
            # Refresh after cleanup
            projects = [p for p in projects if any(v.strip() for v in p.get("queries", {}).values())]
            
            print("\nSelect an option:")
            print("1. Start Schema Generation")
            
            choices = ["1"]
            has_projects = len(projects) > 0
            
            if has_projects:
                print("2. Open Existing Project")
                print("3. Logout")
                choices.extend(["2", "3"])
            else:
                print("2. Logout")
                choices.append("2")
                
            choice = Prompt.ask("Choice", choices=choices)
            
            if choice == "1":
                check_credits_and_start()
            elif choice == "2" and has_projects:
                print("\n[cyan]--- Existing Projects ---[/cyan]")
                for idx, p in enumerate(projects):
                    print(f"{idx + 1}. {p.get('project_name')}")
                    
                proj_choice = Prompt.ask("Select project to open", choices=[str(i+1) for i in range(len(projects))])
                selected_project = projects[int(proj_choice) - 1]
                
                print(f"\n[cyan]Project Name:[/cyan] {selected_project.get('project_name')}")
                project_name = selected_project.get('project_name')
                existing_db_name = selected_project.get('database_name', '')
                
                queries = selected_project.get("queries", {})
                if not queries:
                    print("[italic]No queries saved for this project.[/italic]")
                    input("\nPress Enter to return to main menu...")
                    main(ctx)
                    return

                # Find highest query count dynamically
                highest_query_key = max(queries.keys(), key=lambda k: int(k.replace("query", "")) if k.replace("query", "").isdigit() else -1)
                current_sql = queries[highest_query_key]
                print(f"[cyan]--- Latest Schema ({highest_query_key}) ---[/cyan]")
                print(current_sql)
                print("[cyan]------------------------[/cyan]\n")

                print("Select an option:")
                print("1. Edit Schema")
                print("2. Delete old database & Reprovision")
                print("3. Return to main menu")
                proj_action = Prompt.ask("Choice", choices=["1", "2", "3"])

                if proj_action == "3":
                    main(ctx)
                    return

                if proj_action == "2":
                    ui_data = fetch_user_credits()
                    if ui_data.get('credits', 0) <= 0:
                        print("[red]\nYou have no credits remaining for reprovisioning.[/red]")
                        print("Please upgrade your account or Request more credits to continue.")
                        print("[blue]https://www.aigendb.xyz/pricing[/blue]\n")
                        input("Press Enter to return to main menu...")
                        main(ctx)
                        return

                    # Delete old RDS and reprovision with current schema
                    # Delete old RDS and reprovision with current schema
                    if existing_db_name:
                        print("\n[yellow]Deleting old database...[/yellow]")
                        with console.status("Removing old RDS database..."):
                            delete_database(existing_db_name)

                    print("\n[cyan]Reprovisioning database with current schema...[/cyan]")
                    creds = None
                    while not creds:
                        with console.status("Connecting to database provisioner... (This may take a moment)"):
                            creds = create_database(project_name, current_sql)
                        if not creds:
                            print("[red]\n✗ Database provisioning failed.[/red]")
                            retry = Prompt.ask("What would you like to do?", choices=["retry", "cancel"],
                                               default="retry", show_choices=True, show_default=True)
                            if retry == "cancel":
                                print("[yellow]Reprovisioning cancelled.\n[/yellow]")
                                raise typer.Exit()
                            print("[cyan]Retrying...\n[/cyan]")

                    db_name = creds.get("database_name", "")
                    conn_url = creds.get("connection_url", "")
                    db_pass = creds.get("password", "")

                    with console.status("Saving new credentials..."):
                        save_db_credentials(project_name, db_name, conn_url, db_pass)
                        deduct_user_credit(project_name)

                    project_id = selected_project.get("project_id")
                    url = f"https://www.aigendb.xyz/dashboard/projects/{project_id}"
                    print(f"\n[green]Database is Created successfully, Acess the Database Credentials in you account dashboard : [/green]")
                    print(f"[blue][u]{url}[/u][/blue]\n")
                    typer.launch(url)
                    raise typer.Exit()

                # Open editor
                ui_data = fetch_user_credits()
                if ui_data.get('credits', 0) <= 0:
                    print("[red]\nYou have no credits remaining to edit and reprovision.[/red]")
                    print("Please upgrade your account or Request more credits to continue.")
                    print("[blue]https://www.aigendb.xyz/pricing[/blue]\n")
                    input("Press Enter to return to main menu...")
                    main(ctx)
                    return

                print("\n[cyan]Opening schema editor...[/cyan]")
                print("[italic]Make your changes, save the file, and close the editor to continue.[/italic]")
                new_sql = typer.edit(text=current_sql)

                if not new_sql or new_sql.strip() == current_sql.strip():
                    print("[yellow]No changes detected. Returning to main menu.[/yellow]")
                    input("\nPress Enter to continue...")
                    main(ctx)
                    return

                # Save new query version to Firestore
                with console.status("Saving modified schema..."):
                    append_project_query(project_name, new_sql)
                print("[green]New schema version saved.[/green]")

                # Delete old RDS database
                if existing_db_name:
                    print("\n[yellow]Deleting old database...[/yellow]")
                    with console.status("Removing old RDS database..."):
                        delete_database(existing_db_name)
                
                # Provision new database with retry
                print("\n[cyan]Provisioning new database with updated schema...[/cyan]")
                creds = None
                while not creds:
                    with console.status("Connecting to database provisioner... (This may take a moment)"):
                        creds = create_database(project_name, new_sql)
                    if not creds:
                        print("[red]\n✗ Database provisioning failed.[/red]")
                        retry = Prompt.ask("What would you like to do?", choices=["retry", "cancel"],
                                           default="retry",
                                           show_choices=True,
                                           show_default=True)
                        if retry == "cancel":
                            print("[yellow]Provisioning cancelled. Your schema change is saved safely.\n[/yellow]")
                            raise typer.Exit()
                        print("[cyan]Retrying...\n[/cyan]")

                db_name = creds.get("database_name", "")
                conn_url = creds.get("connection_url", "")
                db_pass = creds.get("password", "")

                with console.status("Saving new credentials..."):
                    save_db_credentials(project_name, db_name, conn_url, db_pass)
                    deduct_user_credit(project_name)

                project_id = selected_project.get("project_id")
                url = f"https://www.aigendb.xyz/dashboard/projects/{project_id}"
                print(f"\n[green]New database provisioned successfully! Access the Database Credentials in your account dashboard : [/green]")
                print(f"[blue][u]{url}[/u][/blue]\n")
                typer.launch(url)
                raise typer.Exit()
            else:
                auth_logout()
                main(ctx)
            return
            
        banner = """
[cyan]    _   ___  ___          ___  ___ [/cyan]
[cyan]   /_\\ |_ _|/ __| ___ _ _|   \\| _ )[/cyan]
[cyan]  / _ \\ | || (_ |/ -_) ' \\ |) | _ \\[/cyan]
[cyan] /_/ \\_\\___|\\___|\\___|_||_|___/___/[/cyan]
"""
        print(banner)
        print("[cyan]Welcome to AIGenDB[/cyan]")
        print("\nSelect an option:")
        print("1. Login to your existing account")
        print("2. Create a new account")
        
        choice = Prompt.ask("Choice", choices=["1", "2"])
        
        if choice == "2":
            print("\nTo create an account please visit:")
            print("[blue]https://www.aigendb.xyz/signup[/blue]\n")
            print("After creating your account return here and login.\n")
            input("Press ENTER to continue...")
            main(ctx)
            return
            
        if choice == "1":
            email = Prompt.ask("Enter your email")
            password = Prompt.ask("Enter your password", password=True)
            
            with console.status("Logging in..."):
                result = auth_login(email, password)
                
            if result.get("success"):
                print("[green]Login successful\n[/green]")
                main(ctx)
            else:
                main(ctx)

@app.command()
def login():
    """Login to AIGenDB"""
    if is_authenticated():
        print("[green]You are already logged in.[/green]")
        return
    email = Prompt.ask("Enter your email")
    password = Prompt.ask("Enter your password", password=True)
    with console.status("Logging in..."):
        result = auth_login(email, password)
    if result.get("success"):
        print("[green]Login successful\n[/green]")

@app.command()
def logout():
    """Logout of AIGenDB"""
    auth_logout()

if __name__ == "__main__":
    app()
