import { StringEnum } from "@earendil-works/pi-ai";
import * as PiCodingAgent from "@earendil-works/pi-coding-agent";
import {
	type AgentSession,
	createAgentSession,
	createExtensionRuntime,
	defineTool,
	type ResourceLoader,
	SessionManager,
	SettingsManager,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import type { TodoObserverConfig } from "./config.ts";
import { normalizeTodos, type TodoNode } from "./types.ts";

const PUBLISH_TOOL = "publish_todo_state";

const SYSTEM_PROMPT = `You are a passive planning observer for another Pi coding session.

Your only job is to maintain a concise hierarchical todo list describing work that the MAIN session intends to complete and work it has completed.

Rules:
- Do not perform the work, inspect files, or give advice.
- Infer todos only from the user's request and strong evidence in the supplied main-session turn.
- Preserve stable todo IDs across updates. Reuse existing IDs whenever a task still represents the same work.
- Use parentId for genuine sub-tasks. Keep the list ordered in the sequence the main session plans to execute it.
- Status must be pending, in_progress, blocked, or completed.
- Mark a task completed only when the turn contains credible evidence that it was completed or verified.
- Mark the task currently being worked on in_progress when that is clear.
- Preserve both unfinished and completed tasks that are not mentioned in the latest turn.
- Completed tasks are historical records. Keep them unless the supplied session context explicitly shows they were invalid, superseded, or intentionally removed.
- Remove tasks only when they were explicitly abandoned, superseded, or are no longer part of the user's request.
- Do not add generic process tasks such as "understand the request", "inspect files", or "respond to user" unless the main session explicitly committed to them as meaningful work.
- Keep todo text short and action-oriented.
- On every observation, call ${PUBLISH_TOOL} exactly once with the complete reconciled list, even when the list is empty.
- Produce no prose outside the tool call.`;

const TodoSchema = Type.Object({
	id: Type.String({ description: "Stable short identifier; preserve it across updates" }),
	parentId: Type.Optional(Type.String({ description: "ID of the parent todo for a sub-task" })),
	text: Type.String({ description: "Short action-oriented todo text" }),
	status: StringEnum(["pending", "in_progress", "blocked", "completed"] as const),
});

function createPublishSchema(maxTodos: number) {
	return Type.Object({
		todos: Type.Array(TodoSchema, {
			description: "The complete current todo tree in display order, represented as a flat parentId list",
			maxItems: maxTodos,
		}),
	});
}

function createObserverResourceLoader(config: TodoObserverConfig): ResourceLoader {
	const additionalInstructions = config.observer.additionalInstructions
		? `\n\nAdditional user configuration:\n${config.observer.additionalInstructions}`
		: "";
	// Source metadata methods were added in Pi 0.83. Keeping them on the
	// structural object is harmless on older Pi versions and supports both APIs.
	return {
		getExtensions: () => ({ extensions: [], errors: [], runtime: createExtensionRuntime() }),
		getSkills: () => ({ skills: [], diagnostics: [] }),
		getPrompts: () => ({ prompts: [], diagnostics: [] }),
		getThemes: () => ({ themes: [], diagnostics: [] }),
		getAgentsFiles: () => ({ agentsFiles: [] }),
		getSystemPrompt: () => SYSTEM_PROMPT + additionalInstructions,
		getSystemPromptSource: () => undefined,
		getAppendSystemPrompt: () => [],
		getAppendSystemPromptSources: () => [],
		extendResources: () => {},
		reload: async () => {},
	} as unknown as ResourceLoader;
}

interface ModelLike {
	provider: string;
	id: string;
}

interface ModelAccess {
	getAvailable(): Promise<ModelLike[]>;
}

interface PiCompatibilityExports {
	ModelRuntime?: {
		create(): Promise<ModelAccess>;
	};
	AuthStorage?: {
		create(): unknown;
	};
	ModelRegistry?: {
		create(authStorage: unknown): ModelAccess;
	};
}

/** Support both Pi's legacy AuthStorage/ModelRegistry SDK and Pi 0.83+ ModelRuntime. */
async function createModelAccess(): Promise<{
	access: ModelAccess;
	sessionOption: { modelRuntime: ModelAccess } | { authStorage: unknown; modelRegistry: ModelAccess };
}> {
	const compat = PiCodingAgent as unknown as PiCompatibilityExports;
	if (compat.ModelRuntime) {
		const modelRuntime = await compat.ModelRuntime.create();
		return { access: modelRuntime, sessionOption: { modelRuntime } };
	}
	if (compat.AuthStorage && compat.ModelRegistry) {
		const authStorage = compat.AuthStorage.create();
		const modelRegistry = compat.ModelRegistry.create(authStorage);
		return { access: modelRegistry, sessionOption: { authStorage, modelRegistry } };
	}
	throw new Error("This Pi version does not expose a supported model/auth SDK");
}

export class TodoObserver {
	private publishedTodos: TodoNode[] | undefined;
	private disposed = false;

	private constructor(
		private readonly session: AgentSession,
		private readonly config: TodoObserverConfig,
	) {}

	static async create(cwd: string, config: TodoObserverConfig): Promise<TodoObserver> {
		const modelAccess = await createModelAccess();
		const availableModels = await modelAccess.access.getAvailable();
		const model = availableModels.find(
			(candidate) => candidate.provider === config.provider && candidate.id === config.model,
		);
		if (!model) {
			throw new Error(
				`Observer model ${config.provider}/${config.model} is unavailable. Authenticate that provider with /login.`,
			);
		}

		let observer: TodoObserver | undefined;
		const publishTool = defineTool({
			name: PUBLISH_TOOL,
			label: "Publish Todo State",
			description: "Publish the complete reconciled todo tree for the observed main Pi session.",
			promptSnippet: "Publish the complete observer todo state",
			parameters: createPublishSchema(config.observer.maxTodos),
			executionMode: "sequential",
			async execute(_toolCallId, params) {
				const todos = normalizeTodos(params.todos as TodoNode[], config.observer.maxTodos);
				if (observer) observer.publishedTodos = todos;
				return {
					content: [{ type: "text" as const, text: `Published ${todos.length} todos.` }],
					details: { todos },
					terminate: true,
				};
			},
		});

		const settingsManager = SettingsManager.inMemory({
			compaction: { enabled: config.observer.compactionEnabled },
			retry: { enabled: config.observer.apiRetries > 0, maxRetries: config.observer.apiRetries },
		});
		const { session } = await createAgentSession({
			cwd,
			model,
			thinkingLevel: config.thinkingLevel,
			...modelAccess.sessionOption,
			resourceLoader: createObserverResourceLoader(config),
			tools: [PUBLISH_TOOL],
			customTools: [publishTool],
			sessionManager: SessionManager.inMemory(cwd),
			settingsManager,
		} as unknown as Parameters<typeof createAgentSession>[0]);

		observer = new TodoObserver(session, config);
		return observer;
	}

	async observe(label: string, mainSessionContext: string, currentTodos: TodoNode[]): Promise<TodoNode[]> {
		if (this.disposed) throw new Error("Observer session is disposed");
		this.publishedTodos = undefined;
		const prompt = `Observation: ${label}

Current todo state:
${JSON.stringify(currentTodos, null, 2)}

New main-session context:
<main_session_context>
${mainSessionContext}
</main_session_context>

Reconcile the complete todo state now and call ${PUBLISH_TOOL}.`;

		await this.session.prompt(prompt);
		for (let attempt = 0; !this.publishedTodos && !this.disposed && attempt < this.config.observer.publishRetries; attempt++) {
			await this.session.prompt(
				`You did not publish a todo state. Call ${PUBLISH_TOOL} now with the complete reconciled list and no prose.`,
			);
		}
		if (!this.publishedTodos) throw new Error("Observer did not publish a todo state");
		return (this.publishedTodos as TodoNode[]).map((todo) => ({ ...todo }));
	}

	dispose(): void {
		if (this.disposed) return;
		this.disposed = true;
		if (this.session.isStreaming) {
			void this.session
				.abort()
				.catch(() => {})
				.finally(() => this.session.dispose());
		} else {
			this.session.dispose();
		}
	}
}