import { readFileSync } from "fs";
import { extname } from "path";

interface MethodData {
	name: string;
	params: string[];
	returns?: string;
	calls: string[];
	raises: string[];
}

interface ClassData {
	name: string;
	methods: MethodData[];
}

export interface ASTFileData {
	exports: string[];
	deps: string[];
	classes: ClassData[];
	functions: MethodData[];
}

const LANGUAGE_MAP: Record<string, string> = {
	".ts": "typescript",
	".tsx": "tsx",
	".js": "javascript",
	".jsx": "javascript",
	".mjs": "javascript",
	".py": "python",
	".go": "go",
	".rs": "rust",
	".java": "java",
	".c": "c",
	".cpp": "cpp",
	".h": "c",
	".rb": "ruby",
};

export async function extractFileAST(
	filePath: string,
): Promise<ASTFileData | null> {
	const ext = extname(filePath).toLowerCase();
	const langName = LANGUAGE_MAP[ext];
	if (!langName) return null;

	try {
		const Parser = require("tree-sitter");
		let grammar: unknown;

		if (langName === "typescript" || langName === "tsx") {
			const ts = require("tree-sitter-typescript");
			grammar = langName === "tsx" ? ts.tsx : ts.typescript;
		} else {
			try {
				const pkg = require(`tree-sitter-${langName}`);
				// Some grammars export the language directly (e.g. python),
				// others export { typescript, tsx } (e.g. typescript)
				grammar = pkg.typescript || pkg.tsx || pkg.go || pkg;
			} catch {
				return null;
			}
		}

		if (!grammar) return null;

		const parser = new Parser();
		parser.setLanguage(grammar);

		const content = readFileSync(filePath, "utf8");
		const tree = parser.parse(content);

		if (langName === "python") {
			return extractPythonData(tree, content);
		}
		if (
			langName === "typescript" ||
			langName === "tsx" ||
			langName === "javascript"
		) {
			return extractTypeScriptData(tree, content);
		}
		if (langName === "go") {
			return extractGoData(tree, content);
		}

		// Fallback for other languages
		const exports = extractExportsFromTree(tree, langName);
		const deps = extractDepsFromTree(tree, langName);
		return { exports, deps, classes: [], functions: [] };
	} catch {
		return null;
	}
}

// ============================================================================
// PYTHON
// ============================================================================

function extractPythonData(tree: Tree, source: string): ASTFileData {
	const classes: ClassData[] = [];
	const exports: string[] = [];
	const deps: string[] = [];
	const functions: MethodData[] = [];

	function visit(node: SyntaxNode) {
		// Imports
		if (
			node.type === "import_statement" ||
			node.type === "import_from_statement"
		) {
			const moduleNode = node.childForFieldName?.("module_name");
			if (moduleNode) {
				deps.push(moduleNode.text);
			} else {
				// import a, b, c
				for (let i = 0; i < node.childCount; i++) {
					const child = node.child(i);
					if (child?.type === "dotted_name" || child?.type === "identifier") {
						deps.push(child.text);
					}
				}
			}
		}

		// Classes
		if (node.type === "class_definition") {
			const nameNode = node.childForFieldName?.("name");
			if (!nameNode) return;
			const className = nameNode.text;
			exports.push(className);

			const methods: MethodData[] = [];
			const body = node.childForFieldName?.("body");
			if (body) {
				for (let i = 0; i < body.childCount; i++) {
					const child = body.child(i);
					if (child?.type === "function_definition") {
						const method = extractPythonMethod(child, source);
						if (method) methods.push(method);
					}
				}
			}
			classes.push({ name: className, methods });
			return; // don't recurse into class body
		}

		// Top-level functions
		if (node.type === "function_definition") {
			const nameNode = node.childForFieldName?.("name");
			if (nameNode) {
				exports.push(nameNode.text);
				const method = extractPythonMethod(node, source);
				if (method) functions.push(method);
			}
		}

		for (let i = 0; i < node.childCount; i++) {
			const child = node.child(i);
			if (child) visit(child);
		}
	}

	visit(tree.rootNode);
	return {
		exports: [...new Set(exports)],
		deps: [...new Set(deps)],
		classes,
		functions,
	};
}

function extractPythonMethod(
	node: SyntaxNode,
	_source: string,
): MethodData | null {
	const nameNode = node.childForFieldName?.("name");
	if (!nameNode) return null;

	const params: string[] = [];
	const parameters = node.childForFieldName?.("parameters");
	if (parameters) {
		for (let i = 0; i < parameters.childCount; i++) {
			const param = parameters.child(i);
			if (param?.type === "identifier" || param?.type === "typed_parameter") {
				params.push(param.text);
			} else if (param?.type === "typed_default_parameter") {
				// name: type = default
				const nameChild = param.childForFieldName?.("name");
				if (nameChild) params.push(nameChild.text);
			}
		}
	}

	// Return type
	let returns: string | undefined;
	const returnType = node.childForFieldName?.("return_type");
	if (returnType) {
		returns = returnType.text.replace(/^->\s*/, "");
	}

	// Calls and raises
	const calls: string[] = [];
	const raises: string[] = [];
	const body = node.childForFieldName?.("body");
	if (body) {
		extractPythonCallsAndRaises(body, calls, raises);
	}

	return {
		name: nameNode.text,
		params,
		returns,
		calls: dedupeCallChains(calls),
		raises: [...new Set(raises)],
	};
}

function dedupeCallChains(calls: string[]): string[] {
	const unique = [...new Set(calls)];
	// Remove shorter calls that are prefixes of longer ones
	return unique.filter(
		(call) =>
			!unique.some(
				(other) =>
					other !== call &&
					(other.startsWith(`${call}.`) || other.startsWith(`${call}(`)),
			),
	);
}

function extractPythonCallsAndRaises(
	node: SyntaxNode,
	calls: string[],
	raises: string[],
) {
	if (node.type === "call") {
		const func = node.childForFieldName?.("function");
		if (func) {
			const callStr = extractCallChain(func);
			if (callStr) calls.push(callStr);
		}
	}
	if (node.type === "raise_statement") {
		const exc = node.child(1);
		if (exc) {
			const excType =
				exc.type === "call" ? exc.childForFieldName?.("function") : exc;
			if (excType) raises.push(excType.text);
		}
	}
	for (let i = 0; i < node.childCount; i++) {
		const child = node.child(i);
		if (child && node.type !== "raise_statement")
			extractPythonCallsAndRaises(child, calls, raises);
	}
}

function extractCallChain(node: SyntaxNode): string | null {
	if (
		node.type === "identifier" ||
		node.type === "attribute" ||
		node.type === "member_expression" ||
		node.type === "property_identifier" ||
		node.type === "type_identifier"
	) {
		return node.text.replace(/\s+/g, " ").trim();
	}
	if (node.type === "call" || node.type === "call_expression") {
		const func = node.childForFieldName?.("function");
		if (func) return extractCallChain(func);
	}
	return null;
}

// ============================================================================
// TYPESCRIPT / JAVASCRIPT
// ============================================================================

function extractTypeScriptData(tree: Tree, source: string): ASTFileData {
	const classes: ClassData[] = [];
	const exports: string[] = [];
	const deps: string[] = [];
	const functions: MethodData[] = [];

	function visit(node: SyntaxNode) {
		// Imports
		if (
			node.type === "import_statement" ||
			node.type === "import_declaration"
		) {
			const sourceNode = node.childForFieldName?.("source");
			if (sourceNode) {
				deps.push(sourceNode.text.slice(1, -1)); // remove quotes
			}
		}
		if (node.type === "call_expression") {
			const func = node.childForFieldName?.("function");
			if (func?.text === "require") {
				const args = node.childForFieldName?.("arguments");
				if (args && args.childCount > 0) {
					const firstArg = args.child(0);
					if (firstArg?.type === "string") {
						deps.push(firstArg.text.slice(1, -1));
					}
				}
			}
		}

		// Classes
		if (node.type === "class_declaration" || node.type === "class") {
			const nameNode = node.childForFieldName?.("name");
			if (!nameNode) return;
			const className = nameNode.text;
			exports.push(className);

			const methods: MethodData[] = [];
			const body = node.childForFieldName?.("body");
			if (body) {
				for (let i = 0; i < body.childCount; i++) {
					const child = body.child(i);
					if (
						child?.type === "method_definition" ||
						child?.type === "function_definition"
					) {
						const method = extractTSMethod(child, source);
						if (method) methods.push(method);
					}
				}
			}
			classes.push({ name: className, methods });
		}

		// Exported functions/consts
		if (
			node.type === "export_statement" ||
			node.type === "export_declaration"
		) {
			const declaration = node.childForFieldName?.("declaration");
			if (declaration) {
				const nameNode = findIdentifier(declaration);
				if (nameNode) {
					exports.push(nameNode.text);
					if (
						declaration.type === "function_declaration" ||
						declaration.type === "function"
					) {
						const method = extractTSMethod(declaration, source);
						if (method) functions.push(method);
					}
				}
			}
		}

		for (let i = 0; i < node.childCount; i++) {
			const child = node.child(i);
			if (child) visit(child);
		}
	}

	visit(tree.rootNode);
	return {
		exports: [...new Set(exports)],
		deps: [...new Set(deps)],
		classes,
		functions,
	};
}

function extractTSMethod(node: SyntaxNode, _source: string): MethodData | null {
	const nameNode = node.childForFieldName?.("name");
	if (!nameNode) return null;

	const params: string[] = [];
	const parameters = node.childForFieldName?.("parameters");
	if (parameters) {
		for (let i = 0; i < parameters.childCount; i++) {
			const param = parameters.child(i);
			if (
				param?.type === "identifier" ||
				param?.type === "required_parameter" ||
				param?.type === "optional_parameter"
			) {
				const name =
					param.childForFieldName?.("pattern") ||
					param.childForFieldName?.("name");
				if (name) {
					const typeAnnotation = param.childForFieldName?.("type");
					if (typeAnnotation) {
						const typeText = typeAnnotation.text.replace(/^:\s*/, "");
						params.push(`${name.text}: ${typeText}`);
					} else {
						params.push(name.text);
					}
				} else {
					params.push(param.text);
				}
			}
		}
	}

	// Return type
	let returns: string | undefined;
	const returnType = node.childForFieldName?.("return_type");
	if (returnType) {
		returns = returnType.text.replace(/^:\s*/, "");
	}

	// Calls and raises
	const calls: string[] = [];
	const raises: string[] = [];
	const body = node.childForFieldName?.("body");
	if (body) {
		extractTSCallsAndThrows(body, calls, raises);
	}

	return {
		name: nameNode.text,
		params,
		returns,
		calls: dedupeCallChains(calls),
		raises: [...new Set(raises)],
	};
}

function extractTSCallsAndThrows(
	node: SyntaxNode,
	calls: string[],
	raises: string[],
) {
	if (node.type === "call_expression") {
		const func = node.childForFieldName?.("function");
		if (func) {
			const callStr = extractCallChain(func);
			if (callStr) calls.push(callStr);
		}
	}
	if (node.type === "throw_statement") {
		const exc = node.child(1);
		if (exc) {
			// Strip 'new ' prefix and extract just the error type
			const text = exc.text.replace(/^new\s+/, "");
			const match = text.match(/^(\w+)/);
			if (match) raises.push(match[1]);
		}
	}
	for (let i = 0; i < node.childCount; i++) {
		const child = node.child(i);
		if (child && node.type !== "throw_statement")
			extractTSCallsAndThrows(child, calls, raises);
	}
}

// ============================================================================
// GO
// ============================================================================

function extractGoData(tree: Tree, _source: string): ASTFileData {
	const exports: string[] = [];
	const deps: string[] = [];

	function visit(node: SyntaxNode) {
		if (node.type === "import_spec") {
			const pathNode = node.childForFieldName?.("path");
			if (pathNode) deps.push(pathNode.text.slice(1, -1));
		}

		if (
			node.type === "function_declaration" ||
			node.type === "method_declaration"
		) {
			const nameNode = node.childForFieldName?.("name");
			if (nameNode && /^[A-Z]/.test(nameNode.text)) {
				exports.push(nameNode.text);
			}
		}

		if (node.type === "type_declaration") {
			const spec = node.childForFieldName?.("spec");
			if (spec) {
				const nameNode = spec.childForFieldName?.("name");
				if (nameNode && /^[A-Z]/.test(nameNode.text)) {
					exports.push(nameNode.text);
				}
			}
		}

		for (let i = 0; i < node.childCount; i++) {
			const child = node.child(i);
			if (child) visit(child);
		}
	}

	visit(tree.rootNode);
	return {
		exports: [...new Set(exports)],
		deps: [...new Set(deps)],
		classes: [],
		functions: [],
	};
}

// ============================================================================
// GENERIC FALLBACK
// ============================================================================

function extractExportsFromTree(tree: Tree, langName: string): string[] {
	const exports: string[] = [];
	const root = tree.rootNode;

	function visit(node: SyntaxNode) {
		if (
			langName === "typescript" ||
			langName === "tsx" ||
			langName === "javascript"
		) {
			if (node.type === "export_statement") {
				const declaration = node.childForFieldName?.("declaration");
				if (declaration) {
					const nameNode = findIdentifier(declaration);
					if (nameNode) exports.push(nameNode.text);
				}
			}
		} else if (langName === "python") {
			if (
				node.type === "function_definition" ||
				node.type === "class_definition"
			) {
				const nameNode = node.childForFieldName?.("name");
				if (nameNode) exports.push(nameNode.text);
			}
		} else if (langName === "go") {
			if (
				node.type === "function_declaration" ||
				node.type === "type_declaration"
			) {
				const nameNode = node.childForFieldName?.("name");
				if (nameNode && /^[A-Z]/.test(nameNode.text)) {
					exports.push(nameNode.text);
				}
			}
		} else if (langName === "rust") {
			if (
				node.type === "function_item" ||
				node.type === "struct_item" ||
				node.type === "enum_item"
			) {
				const nameNode = node.childForFieldName?.("name");
				if (nameNode) exports.push(nameNode.text);
			}
		}

		for (let i = 0; i < node.childCount; i++) {
			const child = node.child(i);
			if (child) visit(child);
		}
	}

	visit(root);
	return [...new Set(exports)];
}

function extractDepsFromTree(tree: Tree, langName: string): string[] {
	const deps: string[] = [];
	const root = tree.rootNode;

	function visit(node: SyntaxNode) {
		if (
			langName === "typescript" ||
			langName === "tsx" ||
			langName === "javascript"
		) {
			if (
				node.type === "import_statement" ||
				node.type === "import_declaration"
			) {
				const source = node.childForFieldName?.("source");
				if (source) deps.push(source.text.slice(1, -1));
			}
			if (node.type === "call_expression") {
				const func = node.childForFieldName?.("function");
				if (func?.text === "require") {
					const args = node.childForFieldName?.("arguments");
					if (args && args.childCount > 0) {
						const firstArg = args.child(0);
						if (firstArg?.type === "string") {
							deps.push(firstArg.text.slice(1, -1));
						}
					}
				}
			}
		} else if (langName === "python") {
			if (
				node.type === "import_statement" ||
				node.type === "import_from_statement"
			) {
				const nameNode = node.childForFieldName?.("name");
				if (nameNode) deps.push(nameNode.text);
			}
		} else if (langName === "go") {
			if (node.type === "import_spec") {
				const pathNode = node.childForFieldName?.("path");
				if (pathNode) deps.push(pathNode.text.slice(1, -1));
			}
		} else if (langName === "rust") {
			if (node.type === "use_declaration") {
				const argument = node.childForFieldName?.("argument");
				if (argument) deps.push(argument.text);
			}
		}

		for (let i = 0; i < node.childCount; i++) {
			const child = node.child(i);
			if (child) visit(child);
		}
	}

	visit(root);
	return [...new Set(deps)];
}

function findIdentifier(node: SyntaxNode): SyntaxNode | null {
	if (
		node.type === "identifier" ||
		node.type === "type_identifier" ||
		node.type === "property_identifier"
	) {
		return node;
	}
	for (let i = 0; i < node.childCount; i++) {
		const child = node.child(i);
		if (child) {
			const found = findIdentifier(child);
			if (found) return found;
		}
	}
	return null;
}

// Type definitions for tree-sitter nodes (simplified)
interface Tree {
	rootNode: SyntaxNode;
}

interface SyntaxNode {
	type: string;
	text: string;
	childCount: number;
	childForFieldName?(name: string): SyntaxNode | null;
	child(index: number): SyntaxNode | null;
}