import { Injectable } from "@nestjs/common";
import { Interval } from "@nestjs/schedule";
import { InjectRepository } from "@nestjs/typeorm";
import * as crypto from "crypto";
import type { Browser, BrowserContext, Page } from "playwright";
import { chromium } from "playwright";
import { Repository } from "typeorm";
import type { ScriptLogger } from "../code-executor/code-executor.service";
import { CodeExecutorService } from "../code-executor/code-executor.service";
import { ExecContextBuilder } from "../code-executor/exec-context.builder";
import { traceStorage } from "../common/trace-context";
import { TraceLogger } from "../common/trace-logger";
import {
  EnvironmentData,
  EnvironmentEntity,
} from "../environment/environment.entity";
import { SessionContextService } from "../session/session-context.service";
import { SessionService } from "../session/session.service";
import { SnippetService } from "../snippet/snippet.service";
import { ScenarioRunLogEntity } from "./scenario-run-log.entity";
import { ScenarioRunStepEntity } from "./scenario-run-step.entity";
import { ScenarioRunEntity } from "./scenario-run.entity";
import { ScenarioStepEntity } from "./scenario-step.entity";
import { ScenarioService } from "./scenario.service";

interface BrowserHandle {
  browser: Browser;
  context: BrowserContext;
  page: Page;
}

@Injectable()
export class ScenarioSchedulerService {
  private readonly logger = new TraceLogger(ScenarioSchedulerService.name);
  private readonly activeRuns = new Set<string>();
  private readonly runBrowsers = new Map<string, BrowserHandle>();
  // Cache credential maps per run (built once when a run starts)
  private readonly runCredentials = new Map<string, Record<string, unknown>>();
  // Cache snippet code map per run (built once when a run starts)
  private readonly runSnippets = new Map<string, Record<string, string>>();
  // Cache environment values per run (built once when a run starts)
  private readonly runEnvironments = new Map<string, EnvironmentData>();
  // Cache scenario-level timeout (seconds) per run
  private readonly runScenarioTimeouts = new Map<string, number | null>();

  private static readonly DEFAULT_SCENARIO_TIMEOUT_SEC = 600;
  private static readonly DEFAULT_STEP_TIMEOUT_SEC = 60;

  constructor(
    @InjectRepository(ScenarioRunEntity)
    private readonly runRepo: Repository<ScenarioRunEntity>,
    @InjectRepository(ScenarioRunStepEntity)
    private readonly runStepRepo: Repository<ScenarioRunStepEntity>,
    @InjectRepository(ScenarioRunLogEntity)
    private readonly runLogRepo: Repository<ScenarioRunLogEntity>,
    @InjectRepository(EnvironmentEntity)
    private readonly environmentRepo: Repository<EnvironmentEntity>,
    private readonly codeExecutor: CodeExecutorService,
    private readonly scenarioService: ScenarioService,
    private readonly snippetService: SnippetService,
    private readonly sessionService: SessionService,
    private readonly sessionContextService: SessionContextService,
  ) {}

  private persistLog(
    runId: string,
    stepRunId: string | null,
    level: "log" | "warn" | "error",
    message: string,
  ): void {
    void this.runLogRepo.save(
      this.runLogRepo.create({ runId, stepRunId, level, message }),
    );
  }

  private stepLogger(stepRunId: string, runId: string): ScriptLogger {
    return (level, msg) => {
      this.logger[level](`StepRun #${stepRunId} script: ${msg}`);
      this.persistLog(runId, stepRunId, level, msg);
    };
  }

  private makeGetStepOutput(
    stepRun: ScenarioRunStepEntity,
  ): (order: number) => Promise<unknown> {
    return async (order: number) => {
      const targetOrder = order < 0 ? stepRun.order + order : order;
      if (targetOrder < 1) return null;
      const sr = await this.runStepRepo.findOne({
        where: { runId: stepRun.runId, order: targetOrder },
      });
      if (!sr?.output) return null;
      try {
        return JSON.parse(sr.output) as unknown;
      } catch {
        return sr.output;
      }
    };
  }

  // ── Job: pick up pending runs and process each to completion ─────────────

  @Interval(1000)
  async pickUpPendingRuns(): Promise<void> {
    const pending = await this.runRepo.find({ where: { status: "pending" } });
    for (const run of pending) {
      if (this.activeRuns.has(run.id)) continue;
      this.activeRuns.add(run.id);
      run.status = "in_progress";
      await this.runRepo.save(run);
      this.logger.log(`Run #${run.id} → in_progress`);
      // Pre-load credential map for the scenario
      const credMap = await this.scenarioService
        .buildCredentialMap(run.scenarioId)
        .catch(() => ({}) as Record<string, unknown>);
      this.runCredentials.set(run.id, credMap);
      // Pre-load snippet map
      const snippetMap = await this.snippetService
        .buildSnippetMap()
        .catch(() => ({}) as Record<string, string>);
      this.runSnippets.set(run.id, snippetMap);
      // Pre-load selected environment data
      const environmentData = await this.environmentRepo
        .findOneBy({ id: run.environmentId })
        .then((env) => env?.data ?? {})
        .catch(() => ({}) as EnvironmentData);
      this.runEnvironments.set(run.id, environmentData);
      // Cache scenario-level timeout for the run
      const scenario = await this.scenarioService
        .findOne(run.scenarioId)
        .catch(() => null);
      this.runScenarioTimeouts.set(run.id, scenario?.timeoutSeconds ?? null);
      const traceId = crypto.randomUUID();
      void traceStorage.run({ traceId }, () =>
        this.processRunToCompletion(run.id),
      );
    }
  }

  private async processRunToCompletion(runId: string): Promise<void> {
    const scenarioTimeoutSec =
      this.runScenarioTimeouts.get(runId) ??
      ScenarioSchedulerService.DEFAULT_SCENARIO_TIMEOUT_SEC;
    const scenarioDeadline = Date.now() + scenarioTimeoutSec * 1000;
    try {
      let stepRun = await this.runStepRepo.findOne({
        where: { runId, status: "pending" },
        relations: ["scenarioStep"],
        order: { order: "ASC" },
      });
      while (stepRun) {
        if (Date.now() >= scenarioDeadline) {
          this.logger.error(
            `Run #${runId}: exceeded scenario timeout of ${scenarioTimeoutSec}s`,
          );
          this.persistLog(
            runId,
            null,
            "error",
            `Scenario timed out after ${scenarioTimeoutSec}s`,
          );
          await this.runRepo.update(runId, { status: "fail" });
          await this.runStepRepo
            .createQueryBuilder()
            .update()
            .set({ status: "cancelled" })
            .where("runId = :runId AND status IN (:...statuses)", {
              runId,
              statuses: ["waiting", "pending", "in_progress"],
            })
            .execute();
          return;
        }
        await this.executeStepRun(stepRun);
        stepRun = await this.runStepRepo.findOne({
          where: { runId, status: "pending" },
          relations: ["scenarioStep"],
          order: { order: "ASC" },
        });
      }
      await this.runRepo
        .createQueryBuilder()
        .update()
        .set({ status: "pass" })
        .where("id = :id AND status = 'in_progress'", { id: runId })
        .execute();
    } catch (err) {
      this.logger.error(
        `Run #${runId}: unexpected error: ${(err as Error).message}`,
      );
      await this.runRepo.update(runId, { status: "fail" });
    } finally {
      this.activeRuns.delete(runId);
      this.runCredentials.delete(runId);
      this.runSnippets.delete(runId);
      this.runEnvironments.delete(runId);
      this.runScenarioTimeouts.delete(runId);
      await this.maybePreserveSession(runId);
    }
  }

  /**
   * If the run was created with saveSession=true, promote the run's browser
   * context to a named explore session instead of closing it.
   */
  private async maybePreserveSession(runId: string): Promise<void> {
    const run = await this.runRepo.findOneBy({ id: runId });
    if (!run?.saveSession) {
      await this.closeBrowserHandle(runId);
      return;
    }

    const handle = this.runBrowsers.get(runId);
    if (!handle) {
      return;
    }

    const sessionName = `run-${runId}`;
    try {
      const cookies = await handle.context.cookies();
      const localStorage = await handle.page
        .evaluate(() => ({ ...window.localStorage }))
        .catch((): Record<string, string> => ({}));
      const token =
        (localStorage["accessToken"] as string | undefined) ??
        (localStorage["token"] as string | undefined) ??
        cookies.find((c) => c.name === "token")?.value ??
        "";

      const session = await this.sessionService.upsert(
        sessionName,
        token,
        cookies,
        localStorage,
      );
      await this.runRepo.update(runId, { sessionId: session.id });
      this.sessionContextService.register(
        sessionName,
        handle.browser,
        handle.context,
        handle.page,
      );
      // Remove from runBrowsers so closeBrowserHandle won't close it
      this.runBrowsers.delete(runId);
      this.logger.log(
        `Run #${runId}: browser preserved as session "${sessionName}" (${session.id})`,
      );
    } catch (err) {
      this.logger.warn(
        `Run #${runId}: failed to preserve session — ${(err as Error).message}`,
      );
      await this.closeBrowserHandle(runId);
    }
  }

  private async executeStepRun(stepRun: ScenarioRunStepEntity): Promise<void> {
    const step = stepRun.scenarioStep as ScenarioStepEntity;

    stepRun.status = "in_progress";
    await this.runStepRepo.save(stepRun);
    this.logger.log(
      `StepRun #${stepRun.id} (run #${stepRun.runId}, step #${step.id} order=${step.order}) → in_progress`,
    );

    try {
      if (!step.execCode) throw new Error("step has no execCode");
      this.codeExecutor.validate(step.execCode);
      const { page, context } = await this.getOrCreateBrowserHandle(
        stepRun.runId,
      );
      const getStepOutput = this.makeGetStepOutput(stepRun);
      const creds = this.runCredentials.get(stepRun.runId);
      const snips = this.runSnippets.get(stepRun.runId);
      const env = this.runEnvironments.get(stepRun.runId);
      const execCtx = new ExecContextBuilder()
        .page(page)
        .browser(context)
        .code(step.execCode)
        .log(this.stepLogger(stepRun.id, stepRun.runId))
        .getStepOutput(getStepOutput)
        .credentials(creds)
        .environment(env)
        .snippets(snips)
        .build();

      const scenarioTimeoutSec = this.runScenarioTimeouts.get(stepRun.runId) ?? null;
      const stepTimeoutSec =
        step.timeoutSeconds ??
        scenarioTimeoutSec ??
        ScenarioSchedulerService.DEFAULT_STEP_TIMEOUT_SEC;
      const stepTimeoutMs = stepTimeoutSec * 1000;

      const timeoutPromise = new Promise<never>((_, reject) =>
        setTimeout(
          () => reject(new Error(`Step timed out after ${stepTimeoutSec}s`)),
          stepTimeoutMs,
        ),
      );
      const { result: execOutput } = await Promise.race([
        this.codeExecutor.execute(execCtx),
        timeoutPromise,
      ]);

      await this.passStepRun(stepRun, null, execOutput);
    } catch (err) {
      const msg = (err as Error).message ?? String(err);
      this.logger.error(`StepRun #${stepRun.id} threw: ${msg}`);
      await this.failStepRun(stepRun, msg);
    }
  }

  // ── Shared browser per run ─────────────────────────────────────────────────

  private async getOrCreateBrowserHandle(
    runId: string,
  ): Promise<BrowserHandle> {
    const existing = this.runBrowsers.get(runId);
    if (existing) return existing;

    const browser = await chromium.launch({
      headless: true,
      executablePath: process.env.PLAYWRIGHT_CHROMIUM_EXECUTABLE_PATH,
    });
    const context = await browser.newContext();
    const page = await context.newPage();

    const handle: BrowserHandle = { browser, context, page };
    this.runBrowsers.set(runId, handle);
    this.logger.log(`Run #${runId}: browser created`);
    return handle;
  }

  private async closeBrowserHandle(runId: string): Promise<void> {
    const handle = this.runBrowsers.get(runId);
    if (!handle) return;
    this.runBrowsers.delete(runId);
    try {
      await handle.browser.close();
      this.logger.log(`Run #${runId}: browser closed`);
    } catch (err) {
      this.logger.warn(
        `Run #${runId}: error closing browser: ${(err as Error).message}`,
      );
    }
  }

  // ── Pass / fail helpers ────────────────────────────────────────────────────

  private async passStepRun(
    stepRun: ScenarioRunStepEntity,
    description: string | null,
    output: unknown = null,
  ): Promise<void> {
    stepRun.status = "pass";
    stepRun.description = description;
    stepRun.output =
      output !== null && output !== undefined ? JSON.stringify(output) : null;
    await this.runStepRepo.save(stepRun);
    this.logger.log(`StepRun #${stepRun.id} → pass`);

    // Find the next waiting step in this run (next by order)
    const nextStep = await this.runStepRepo.findOne({
      where: { runId: stepRun.runId, status: "waiting" },
      order: { order: "ASC" },
    });

    if (nextStep) {
      nextStep.status = "pending";
      await this.runStepRepo.save(nextStep);
    } else {
      // No more waiting steps — check if any are still in_progress/pending (shouldn't be, but guard anyway)
      const remaining = await this.runStepRepo.count({
        where: [
          { runId: stepRun.runId, status: "pending" },
          { runId: stepRun.runId, status: "in_progress" },
          { runId: stepRun.runId, status: "waiting" },
        ],
      });
      if (remaining === 0) {
        // Don't close the browser here — maybePreserveSession will handle it
        // (either preserving it if saveSession=true, or closing if false)
        await this.runRepo.update(stepRun.runId, { status: "pass" });
        this.logger.log(`Run #${stepRun.runId} → pass (all steps passed)`);
      }
    }
  }

  private async failStepRun(
    stepRun: ScenarioRunStepEntity,
    description: string,
  ): Promise<void> {
    stepRun.status = "fail";
    stepRun.description = description;
    await this.runStepRepo.save(stepRun);
    this.logger.log(`StepRun #${stepRun.id} → fail: ${description}`);

    // Cancel all remaining waiting/pending step runs in this run
    await this.runStepRepo
      .createQueryBuilder()
      .update()
      .set({ status: "cancelled" })
      .where("runId = :runId AND status IN (:...statuses)", {
        runId: stepRun.runId,
        statuses: ["waiting", "pending"],
      })
      .execute();

    this.logger.log(`Run #${stepRun.runId}: remaining steps cancelled`);

    // Don't close the browser here — maybePreserveSession (called from
    // processRunToCompletion's finally block) will either preserve it as a
    // session (saveSession=true) or close it (saveSession=false).
    await this.runRepo.update(stepRun.runId, { status: "fail" });
    this.logger.log(`Run #${stepRun.runId} → fail`);
  }
}