import { execa } from "execa";
import puppeteer, { Browser, Page } from "puppeteer";
import { mkdir, writeFile } from "fs/promises";
import path from "path";
import stripAnsi from "strip-ansi";
import os from "os";
import inquirer from "inquirer";

const SESSION_DIR = path.join(os.homedir(), ".miso", "sessions");

interface Session {
  alias: string;
  frontdoorUrl: string;
  instanceUrl: string;
  sid: string;
  cookies: any[];
  auraToken: string;
  savedAt: number;
}

interface OrgInfo {
  alias: string;
  username: string;
  instanceUrl: string;
  name?: string;
  isDevHub?: boolean;
  isScratch?: boolean;
  isSandbox?: boolean;
  connectedStatus?: string;
}

async function listOrgs(): Promise<OrgInfo[]> {
  const { stdout } = await execa("sf", ["org", "list", "--json"]);
  const data = JSON.parse(stdout);
  // Combine all orgs from different categories
  const allOrgs = [
    ...(data.result.other || []),
    ...(data.result.sandboxes || []),
    ...(data.result.nonScratchOrgs || []),
    ...(data.result.devHubs || []),
    ...(data.result.scratchOrgs || [])
  ];

  // Filter out orgs with AuthDecryptError
  const validOrgs = allOrgs.filter(org => org.connectedStatus !== "AuthDecryptError");

  // Map to our OrgInfo interface
  return validOrgs.map(org => ({
    alias: org.alias,
    username: org.username,
    instanceUrl: org.instanceUrl,
    name: org.name,
    isDevHub: org.isDevHub,
    isScratch: org.isScratch,
    isSandbox: org.isSandbox,
    connectedStatus: org.connectedStatus
  }));
}

export async function selectOrg(): Promise<string> {
  const orgs = await listOrgs();

  if (orgs.length === 0) {
    throw new Error("No Salesforce orgs found. Please run 'sf org login' first.");
  }

  const { selectedAlias } = await inquirer.prompt([
    {
      type: "list",
      name: "selectedAlias",
      message: "Select a Salesforce org:",
      choices: orgs.map((org) => {
        let name = `${org.alias} (${org.username})`;
        if (org.isDevHub) name += ' [DevHub]';
        if (org.isScratch) name += ' [Scratch]';
        if (org.isSandbox) name += ' [Sandbox]';
        if (org.name) name += ` - ${org.name}`;
        return {
          name,
          value: org.alias,
        };
      }),
    },
  ]);

  return selectedAlias;
}

async function getFrontdoorUrl(
  alias?: string
): Promise<{ frontdoorUrl: string; sid: string }> {
  const selectedAlias = alias;

  const args = ["org", "open", "-r"];
  if (selectedAlias) {
    args.push("-o", selectedAlias);
  }

  const { stdout } = await execa("sf", args);
  const match = stripAnsi(stdout).match(
    /https:\/\/[^ ]+frontdoor\.jsp\?sid=([^\s]+)/
  );

  if (!match) {
    throw new Error(`Failed to get frontdoor URL: ${stdout}`);
  }

  return {
    frontdoorUrl: match[0],
    sid: match[1]!,
  };
}

async function initializeBrowser(headless: boolean = true): Promise<Browser> {
  return await puppeteer.launch({
    headless,
    args: ["--no-sandbox", "--disable-setuid-sandbox"],
  });
}

async function waitForAuraToken(
  page: Page,
  timeout: number = 10000
): Promise<string | null> {
  try {
    await page.waitForFunction(
      () => {
        return (
          // @ts-ignore
          window.Aura?.initConfig?.token ||
          // @ts-ignore
          window.aura?.services?.$client$?.$_token$
        );
      },
      { timeout }
    );

    const token = await page.evaluate(() => {
      // @ts-ignore
      return window.Aura.initConfig?.token || window.aura?.services?.$client$?.$_token$;
    });

    return token || null;
  } catch (error) {
    console.log("Failed to get Aura token:", error);
    return null;
  }
}

async function getAuraTokenWithRetry(
  page: Page,
  maxRetries: number = 3
): Promise<string> {
  for (let attempt = 1; attempt <= maxRetries; attempt++) {
    console.log(`Attempt ${attempt}/${maxRetries} to get Aura token...`);

    const token = await waitForAuraToken(page);
    if (token) {
      return token;
    }

    if (attempt < maxRetries) {
      console.log("Waiting before retry...");
      await new Promise((resolve) => setTimeout(resolve, 2000));
    }
  }

  throw new Error("Failed to get Aura token after all attempts");
}

async function navigateToLightning(
  page: Page,
  frontdoorUrl: string,
  instanceUrl: string
): Promise<void> {
  await page.goto(frontdoorUrl, { waitUntil: "networkidle2" });
  const lightningAppUrl = instanceUrl + "/lightning/page/home";
  await page.goto(lightningAppUrl, { waitUntil: "networkidle2" });
}

async function saveSession(session: Session): Promise<void> {
  await mkdir(SESSION_DIR, { recursive: true });
  await writeFile(
    path.join(SESSION_DIR, `${session.alias}.json`),
    JSON.stringify(session, null, 2),
    "utf-8"
  );
}

export async function getSfSession(alias?: string): Promise<Session> {

  // Step 1: Get frontdoor URL (now with org selection if needed)
  const { frontdoorUrl, sid } = await getFrontdoorUrl(alias);
  const instanceUrl = new URL(frontdoorUrl).origin;

  // Step 2: Initialize browser and try to get token
  let browser = await initializeBrowser(true);
  let page = await browser.newPage();

  try {
    await navigateToLightning(page, frontdoorUrl, instanceUrl);

    // Step 3: Try to get Aura token
    let auraToken: string;
    try {
      auraToken = await getAuraTokenWithRetry(page, 1);
    } catch (error) {
      console.log(
        "Failed to get token in headless mode, retrying in non-headless mode..."
      );

      // Close headless browser and try with non-headless
      await browser.close();
      browser = await initializeBrowser(false);
      page = await browser.newPage();

      await navigateToLightning(page, frontdoorUrl, instanceUrl);
      auraToken = await getAuraTokenWithRetry(page, 1);
    }

    // Step 4: Get cookies and instance URL
    const cookies = await page.cookies();
    const newInstanceUrl = await page.evaluate(() => window.location.origin);

    // Step 5: Create and save session
    const session: Session = {
      alias: alias || (await selectOrg()), // Use selected org alias
      frontdoorUrl,
      instanceUrl: newInstanceUrl,
      sid,
      cookies,
      auraToken,
      savedAt: Date.now(),
    };

    await saveSession(session);
    console.log("✅ Login success!");
    return session;
  } finally {
    await browser.close();
  }
}