import { z } from "zod";
import {
  AgentSchema,
  AlertSchema,
  CatalystSchema,
  CompanySchema,
  EarningsScheduleSchema,
  ExportRecordSchema,
  FilingSchema,
  HoldingSchema,
  MemoAnnotationSchema,
  MemoCitationSchema,
  MemoSectionReviewSchema,
  MemoSectionSchema,
  ModelRowSchema,
  RiskSchema,
  WorkspaceSectionSchema,
  type ClientSettings,
  type RpcMethod,
  type RpcRequestMap,
  type RpcResponseMap,
  type ServerSettings,
} from "./rpc.js";

const nonEmptyString = z.string().trim().min(1);
const idString = nonEmptyString;
const tickerString = z.string().trim().min(1).max(16);
const nonNegativeIndex = z.number().int().min(0);
const unknownRecord = z.record(z.unknown());

export const ClientSettingsSchema = z.object({
  theme: z.enum(["light", "dark", "system"]),
  density: z.enum(["comfortable", "compact", "dense"]),
  sidebarWidth: z.number().int().min(160).max(520),
  navCollapsed: z.record(z.boolean()),
  keybindings: z.record(z.string()),
});

export const ServerSettingsSchema = z.object({
  agentConfigs: z.record(z.unknown()),
  dataSources: z.record(z.boolean()),
  exportPipelines: z.record(z.unknown()),
});

const RiskInputSchema = RiskSchema.omit({ id: true, companyId: true });
const SettingsUpdateSchema = z.discriminatedUnion("scope", [
  z.object({ scope: z.literal("client"), changes: ClientSettingsSchema.partial() }),
  z.object({ scope: z.literal("server"), changes: ServerSettingsSchema.partial() }),
]);

export const RpcRequestSchemas = {
  "portfolio.get": z.undefined(),
  "portfolio.addHolding": z.object({ ticker: tickerString }),
  "portfolio.removeHolding": z.object({ ticker: tickerString }),
  "company.get": z.object({ companyId: idString }),
  "company.search": z.object({ query: nonEmptyString }),
  "company.setActive": z.object({ companyId: idString }),
  "workspace.getSection": z.object({ companyId: idString, section: nonEmptyString }),
  "workspace.listSources": z.object({ companyId: idString }),
  "catalyst.list": z.object({ companyId: idString }),
  "alert.list": z.object({ companyId: idString.optional(), since: z.string().optional() }),
  "risk.list": z.object({ companyId: idString }),
  "risk.add": z.object({ companyId: idString, risk: RiskInputSchema }),
  "earnings.getSchedule": z.object({ companyId: idString }),
  "filing.list": z.object({ companyId: idString, since: z.string().optional() }),
  "model.get": z.object({ companyId: idString, tab: nonEmptyString }),
  "model.updateCell": z.object({
    companyId: idString,
    tab: nonEmptyString,
    row: nonNegativeIndex,
    col: nonNegativeIndex,
    value: z.string(),
  }),
  "model.runScenario": z.object({
    companyId: idString,
    scenario: nonEmptyString,
    overrides: z.record(z.string()),
  }),
  "memo.get": z.object({ companyId: idString }),
  "memo.updateSection": z.object({
    companyId: idString,
    sectionId: idString,
    title: nonEmptyString.optional(),
    content: nonEmptyString,
  }),
  "memo.addAnnotation": z.object({
    companyId: idString,
    sectionId: idString,
    kind: z.enum(["highlight", "comment", "strike"]),
    selectedText: nonEmptyString,
    comment: nonEmptyString.optional(),
  }),
  "memo.resolveAnnotation": z.object({ companyId: idString, annotationId: idString }),
  "memo.updateSectionReview": z.object({
    companyId: idString,
    sectionId: idString,
    status: z.enum(["pending", "in_review", "approved", "changes_requested"]),
  }),
  "memo.acceptEdit": z.object({ companyId: idString, editId: idString }),
  "memo.rejectEdit": z.object({ companyId: idString, editId: idString, reason: z.string().optional() }),
  "agent.list": z.object({ companyId: idString.optional() }),
  "agent.start": z.object({ agentId: idString, companyId: idString }),
  "agent.pause": z.object({ agentId: idString }),
  "agent.restart": z.object({ agentId: idString }),
  "agent.chat": z.object({ agentId: idString, message: nonEmptyString }),
  "agent.configure": z.object({ agentId: idString, config: unknownRecord }),
  "agent.getTrace": z.object({ agentId: idString, runId: idString }),
  "agent.runPipeline": z.object({
    companyId: idString,
    pipeline: z.enum(["research", "competitive", "cross-cutting"]),
  }),
  "validation.run": z.object({ companyId: idString, agentType: z.enum(["sv", "qa", "rt", "all"]).optional() }),
  "validation.getStatus": z.object({ companyId: idString, sectionId: idString.optional() }),
  "export.list": z.object({ companyId: idString.optional() }),
  "export.create": z.object({ type: z.enum(["pdf", "excel", "ppt"]), companyId: idString, options: unknownRecord.optional() }),
  "export.download": z.object({ exportId: idString }),
  "settings.get": z.object({ scope: z.enum(["client", "server"]) }),
  "settings.update": SettingsUpdateSchema,
} satisfies Record<RpcMethod, z.ZodTypeAny>;

const SettingsResponseSchema = z.union([ClientSettingsSchema, ServerSettingsSchema]);
const ValidationIssueSchema = z.object({
  severity: z.string(),
  message: z.string(),
  suggestion: z.string().optional(),
});
const ValidationResultSchema = z.object({
  passed: z.boolean(),
  confidence: z.string(),
  issues: z.array(ValidationIssueSchema),
  notes: z.string(),
  timestamp: z.string(),
});

export const RpcResponseSchemas = {
  "portfolio.get": z.object({ id: z.string(), name: z.string(), holdings: z.array(HoldingSchema), activeCompanyId: z.string() }),
  "portfolio.addHolding": z.object({ holding: HoldingSchema }),
  "portfolio.removeHolding": z.object({ ok: z.boolean() }),
  "company.get": z.object({ company: CompanySchema }),
  "company.search": z.object({ results: z.array(z.object({ ticker: z.string(), name: z.string(), sector: z.string() })) }),
  "company.setActive": z.object({ ok: z.boolean() }),
  "workspace.getSection": z.object({ content: WorkspaceSectionSchema, validationState: z.string() }),
  "workspace.listSources": z.object({ sources: z.array(z.object({ type: z.string(), title: z.string(), metadata: z.string() })) }),
  "catalyst.list": z.object({ catalysts: z.array(CatalystSchema) }),
  "alert.list": z.object({ alerts: z.array(AlertSchema) }),
  "risk.list": z.object({ risks: z.array(RiskSchema) }),
  "risk.add": z.object({ risk: RiskSchema }),
  "earnings.getSchedule": z.object({ schedule: z.array(EarningsScheduleSchema) }),
  "filing.list": z.object({ filings: z.array(FilingSchema) }),
  "model.get": z.object({ headers: z.array(z.string()), rows: z.array(ModelRowSchema) }),
  "model.updateCell": z.object({ ok: z.boolean(), affectedCells: z.array(z.string()) }),
  "model.runScenario": z.object({ headers: z.array(z.string()), rows: z.array(ModelRowSchema) }),
  "memo.get": z.object({
    status: z.enum(["draft", "review", "final"]),
    sections: z.array(MemoSectionSchema),
    citations: z.array(MemoCitationSchema),
    annotations: z.array(MemoAnnotationSchema),
    sectionReviews: z.array(MemoSectionReviewSchema),
  }),
  "memo.updateSection": z.object({
    section: MemoSectionSchema,
    status: z.enum(["draft", "review", "final"]),
    savedAt: z.string(),
  }),
  "memo.addAnnotation": z.object({ annotation: MemoAnnotationSchema }),
  "memo.resolveAnnotation": z.object({ annotation: MemoAnnotationSchema }),
  "memo.updateSectionReview": z.object({ review: MemoSectionReviewSchema }),
  "memo.acceptEdit": z.object({ ok: z.boolean() }),
  "memo.rejectEdit": z.object({ ok: z.boolean() }),
  "agent.list": z.object({ agents: z.array(AgentSchema) }),
  "agent.start": z.object({ runId: z.string() }),
  "agent.pause": z.object({ ok: z.boolean() }),
  "agent.restart": z.object({ runId: z.string() }),
  "agent.chat": z.object({ response: z.string() }),
  "agent.configure": z.object({ ok: z.boolean() }),
  "agent.getTrace": z.object({ steps: z.array(z.object({ step: z.number(), label: z.string(), detail: z.string() })) }),
  "agent.runPipeline": z.object({ runIds: z.array(z.string()) }),
  "validation.run": z.object({
    sourceVerification: ValidationResultSchema.optional(),
    modelQA: ValidationResultSchema.optional(),
    redTeam: ValidationResultSchema.optional(),
  }),
  "validation.getStatus": z.object({
    validationState: z.enum(["verified", "flagged", "unverified", "failed"]),
    lastValidated: z.string().optional(),
  }),
  "export.list": z.object({ exports: z.array(ExportRecordSchema) }),
  "export.create": z.object({ exportId: z.string() }),
  "export.download": z.object({ data: z.instanceof(ArrayBuffer) }),
  "settings.get": z.object({ settings: SettingsResponseSchema }),
  "settings.update": z.object({ ok: z.boolean() }),
} satisfies Record<RpcMethod, z.ZodTypeAny>;

export const AgentProgressEventSchema = z.object({
  type: z.literal("agent.progress"),
  data: z.object({ runId: z.string(), agentId: z.string(), companyId: z.string(), progress: z.number(), action: z.string(), timestamp: z.string() }),
});
export const AgentCompletedEventSchema = z.object({
  type: z.literal("agent.completed"),
  data: z.object({ runId: z.string(), agentId: z.string(), companyId: z.string(), output: z.unknown(), duration: z.number(), timestamp: z.string() }),
});
export const AgentFailedEventSchema = z.object({
  type: z.literal("agent.failed"),
  data: z.object({ runId: z.string(), agentId: z.string(), companyId: z.string(), error: z.string(), timestamp: z.string() }),
});
export const AgentStartedEventSchema = z.object({
  type: z.literal("agent.started"),
  data: z.object({ runId: z.string(), agentId: z.string(), companyId: z.string(), pipeline: z.string().optional(), timestamp: z.string() }),
});
export const AgentStreamingEventSchema = z.object({
  type: z.literal("agent.streaming"),
  data: z.object({ runId: z.string(), agentId: z.string(), companyId: z.string(), chunk: z.string(), timestamp: z.string() }),
});
export const ValidationUpdatedEventSchema = z.object({
  type: z.literal("validation.updated"),
  data: z.object({
    companyId: z.string(),
    sectionId: z.string().optional(),
    validationState: z.enum(["verified", "flagged", "unverified", "failed"]),
    agentId: z.string(),
    notes: z.string().optional(),
    timestamp: z.string(),
  }),
});
export const MemoUpdatedEventSchema = z.object({
  type: z.literal("memo.updated"),
  data: z.object({ companyId: z.string(), sectionId: z.string(), content: z.string(), updatedAt: z.string() }),
});
export const ModelUpdatedEventSchema = z.object({
  type: z.literal("model.updated"),
  data: z.object({ companyId: z.string(), tab: z.string(), cell: z.object({ row: z.number(), col: z.number() }).optional(), timestamp: z.string() }),
});

export const ServerEventSchema = z.discriminatedUnion("type", [
  AgentProgressEventSchema,
  AgentCompletedEventSchema,
  AgentFailedEventSchema,
  AgentStartedEventSchema,
  AgentStreamingEventSchema,
  ValidationUpdatedEventSchema,
  MemoUpdatedEventSchema,
  ModelUpdatedEventSchema,
]);

export const ServerEventTypeSchema = z.enum([
  "agent.progress",
  "agent.completed",
  "agent.failed",
  "agent.started",
  "agent.streaming",
  "validation.updated",
  "memo.updated",
  "model.updated",
]);

export function isRpcMethod(method: unknown): method is RpcMethod {
  return typeof method === "string" && method in RpcRequestSchemas;
}

export function parseRpcRequest<T extends RpcMethod>(method: T, payload: unknown): RpcRequestMap[T] {
  return RpcRequestSchemas[method].parse(payload) as RpcRequestMap[T];
}

export function parseRpcResponse<T extends RpcMethod>(method: T, data: unknown): RpcResponseMap[T] {
  return RpcResponseSchemas[method].parse(data) as RpcResponseMap[T];
}

export type SettingsUpdatePayload =
  | { scope: "client"; changes: Partial<ClientSettings> }
  | { scope: "server"; changes: Partial<ServerSettings> };