################################################################################
# MIT License

# Copyright (c) 2025 Advanced Micro Devices, Inc. All Rights Reserved.

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
################################################################################

import json
import logging
import os

from intelliperf.core.llm import LLM
from intelliperf.formulas.formula_base import (
	Formula_Base,
	Result,
	filter_json_field,
	get_kernel_name,
)
from intelliperf.utils.env import get_llm_api_key


class atomic_contention(Formula_Base):
	def __init__(
		self,
		name: str,
		build_command: list,
		instrument_command: list,
		project_directory: str,
		app_cmd: list,
		top_n: int,
		model: str = "gpt-4o",
		provider: str = "openai",
		in_place: bool = False,
		**kwargs,
	):
		super().__init__(
			name,
			build_command,
			instrument_command,
			project_directory,
			app_cmd,
			top_n,
			model,
			provider,
			in_place,
			**kwargs,
		)

		self._instrumentation_results = None
		self.current_kernel = None
		self.current_args = None
		self.current_kernel_signature = None
		self.kernel_to_optimize = None
		self.optimization_report = None
		self.bottleneck_report = None
		self.current_summary = None
		self.previous_source_code = None
		self.success = False
		self.optimization_attempts = []  # Track optimization strategies attempted
		self.iteration_count = 0  # Track current iteration

	def profile_pass(self) -> Result:
		"""
		Profile the application using guided-tuning and collect atomic contention data

		Returns:
		    Result: DataFrame containing the performance report card
		"""
		super().profile_pass()

		# Reset optimization state for new optimization cycle
		self.reset_optimization_state()

		# Log profiling results
		if hasattr(self, "_initial_profiler_results") and self._initial_profiler_results:
			self.get_logger().record(
				"profile_pass_complete", {"profiler_results": self._initial_profiler_results, "top_n": self.top_n}
			)

	def instrument_pass(self) -> Result:
		"""
		Instrument the application, targeting the kernels with the highest atomic contention data

		Returns:
		    Result: Instrumentation data containing the kernel name, arguments, lines, and file path as dict
		"""
		super().instrument_pass()

		# Log instrumentation completion
		self.get_logger().record(
			"instrument_pass_complete", {"success": True, "note": "Instrumentation pass completed via parent class"}
		)

		return Result(
			success=False,
			asset=self._instrumentation_results,
			error_report="Instrumentation pass not implemented for atomic contention.",
		)

	def build_pass(self, validate_build_result=True) -> Result:
		"""
		Build the application and store the summary.

		Args:
		    validate_build_result (bool): Whether to validate the build result

		Returns:
		    Result: Build status and the output file path
		"""
		result = super().build(validate_build_result=validate_build_result)
		if not result:
			logging.debug(f"Setting current summary to: {result.error_report}")
			self.current_summary = result.error_report
		return result

	def optimize_pass(self, temperature: float = 0.0, max_tokens: int = 3000, target_kernel: str = None) -> Result:
		"""
		Optimize the kernel to remove atomic contention via OpenAI API

		Args:
		    temperature (float): Sampling temperature for OpenAI API
		    max_tokens (int): Maximum tokens for OpenAI API

		Returns:
		    Result: Optimized kernel as a file path
		"""
		super().optimize_pass()
		llm_key = get_llm_api_key()

		system_prompt = (
			"You are a skilled GPU HIP programmer. Given a kernel,"
			" you will optimize it to remove atomic contention as much as possible"
			" and provide a correct performant implementation. Do not modify"
			" the kernel signature. Do not touch any other code, licenses, copyrights, or comments in the file."
			" If you remove the copyright, your solution will be rejected."
			" Do not include any markdown code blocks or text other than the code."
		)

		model = self.model
		provider = self.provider
		llm = LLM(
			api_key=llm_key,
			system_prompt=system_prompt,
			model=model,
			provider=provider,
			logger=self.get_logger(),  # Add logger here
		)

		kernel = None
		kernel_file = None

		if self._instrumentation_results is None:
			# Get the file from the results
			field = "atomics"
			subfield = "atomic_lat"
			# Average atomic latency in cycles measured experimentally
			average_atomic_lat = 1000
			filtered_report_card = filter_json_field(
				self._initial_profiler_results,
				field=field,
				subfield=subfield,
				comparison_func=lambda x: x > average_atomic_lat,
				target_kernel=target_kernel,
			)

			if len(filtered_report_card) == 0:
				return Result(success=False, error_report="No atomic contention found.")

			logging.debug(f"Filtered Report Card:\n{json.dumps(filtered_report_card, indent=4)}")

			kernel = filtered_report_card[0]["kernel"]
			self._parse_kernel_signature(kernel)
			files = filtered_report_card[0]["source"]["files"]
			kernel_name = get_kernel_name(kernel)
			kernel_file = None
			unoptimized_file_content = None
			for file in files:
				project_dir = os.path.abspath(self._application.get_project_directory())
				file_path = os.path.abspath(file)
				isfile_in_project = os.path.commonpath([project_dir, file_path]) == project_dir
				if os.path.exists(file) and isfile_in_project:
					with open(file, "r") as f:
						unoptimized_file_content = f.read()
						if kernel_name in unoptimized_file_content:
							kernel_file = file
							break
			if kernel_file is None:
				return Result(success=False, error_report="Kernel file not found.")

			user_prompt = (
				f"OPTIMIZATION ITERATION: You are optimizing kernel '{kernel}' to reduce atomic contention.\n\n"
				f"SOURCE CODE:\n{unoptimized_file_content}\n\n"
			)

			# Add iteration context if this is not the first attempt
			if self.current_summary is not None:
				user_prompt += (
					f"PREVIOUS OPTIMIZATION RESULTS:\n{self.current_summary}\n\n"
					"The previous optimization approach did not improve atomic contention. "
					"Try a different strategy.\n\n"
				)

				# Add optimization history context
				history_context = self.get_optimization_history_context()
				user_prompt += f"{history_context}\n\n"

				# Add diff information to show what was changed
				cur_diff = self.compute_diff([kernel_file])
				if cur_diff:
					user_prompt += f"CHANGES FROM ORIGINAL CODE:\n{cur_diff}\n\n"

			user_prompt += (
				"OBJECTIVE: Reduce atomic contention while maintaining or improving runtime performance.\n"
				"CONSTRAINTS: Preserve kernel signature and program semantics.\n"
				"Return only the optimized code without any markdown or explanatory text."
			)

			self.previous_source_code = unoptimized_file_content

			if self.current_args:
				self.bottleneck_report = (
					f"Atomic Contention Detection: IntelliPerf identified high atomic contention in kernel "
					f"`{self.current_kernel}` with arguments `{self.current_args}`. Atomic contention occurs when multiple threads "
					f"compete for the same atomic operations, causing serialization and increased latency."
				)
			else:
				self.bottleneck_report = (
					f"Atomic Contention Detection: IntelliPerf identified high atomic contention in kernel "
					f"`{self.current_kernel}`. Atomic contention occurs when multiple threads "
					f"compete for the same atomic operations, causing serialization and increased latency."
				)

		else:
			pass

		if kernel is None:
			return Result(success=False, error_report="Failed to extract the kernel name.")
		if kernel_file is None:
			return Result(success=False, error_report="Failed to extract the kernel file path.")

		self.current_kernel_files = [kernel_file]

		logging.debug(f"LLM prompt: {user_prompt}")
		logging.debug(f"System prompt: {system_prompt}")

		try:
			record_meta = f"Iteration {self.iteration_count + 1}"
			optimized_file_content = llm.ask(user_prompt, record_meta=record_meta).strip()
			optimized_file_content = self.postprocess_llm_code(optimized_file_content)

			# Track this optimization attempt
			self.iteration_count += 1

			# Store attempt info (will be updated with results after validation)
			self.optimization_attempts.append(
				{
					"iteration": self.iteration_count,
					"code": optimized_file_content,
					"success": None,  # Will be updated after performance validation
					"improvement": 0.0,  # Will be updated after performance validation
				}
			)

			# Log successful optimization
			self.get_logger().record(
				"optimization_success",
				{
					"optimized_code_length": len(optimized_file_content),
					"kernel_file": kernel_file,
					"iteration": self.iteration_count,
				},
			)

			with open(kernel_file, "w") as f:
				f.write(optimized_file_content)
			logging.debug(f"Optimized file content: {optimized_file_content}")
			return Result(
				success=True,
				asset={
					"optimized_code_path": kernel_file,
					"optimized_code_string": optimized_file_content,
				},
			)
		except Exception as e:
			error_msg = f"An unexpected error occurred - {str(e)}"
			self.get_logger().record("optimization_error", {"error": error_msg, "error_type": type(e).__name__})
			logging.error(error_msg)
			return Result(success=False, error_report=error_msg)

	def compiler_pass(self) -> Result:
		"""
		Compile the application

		Returns:
		    Result: Compilation status and the output file path
		"""
		return super().compile_pass()

	def correctness_validation_pass(self, accordo_absolute_tolerance: float = 1e-6) -> Result:
		"""
		Validate the optimized kernel by comparing the output with the reference kernel

		Args:
		    accordo_absolute_tolerance (float): Absolute tolerance for Accordo
		Returns:
		    Result: Validation status
		"""
		result = super().correctness_validation_pass(self.current_kernel, self.current_args, accordo_absolute_tolerance)
		if not result:
			logging.info(f"Setting current summary to: {result.error_report}")
			self.current_summary = result.error_report
		return result

	def performance_validation_pass(self) -> Result:
		"""
		Validate the optimized kernel by comparing the output with the reference kernel

		Returns:
		    Result: Validation status
		"""
		unoptimized_results = filter_json_field(
			self._initial_profiler_results,
			field="kernel",
			comparison_func=lambda x: x == self.current_kernel_signature,
		)

		unoptimized_time = unoptimized_results[0]["durations"]["ns"]

		field = "atomics"
		subfield = "atomic_lat"
		unoptimized_metric = unoptimized_results[0][field][subfield]

		# Profile the optimized application
		self._optimization_results = self._application.profile(top_n=self.top_n)

		optimized_results = filter_json_field(
			self._optimization_results,
			field="kernel",
			comparison_func=lambda x: x == self.current_kernel_signature,
		)

		optimized_time = optimized_results[0]["durations"]["ns"]
		optimized_metric = optimized_results[0][field][subfield]

		success = optimized_metric < unoptimized_metric
		speedup = unoptimized_time / optimized_time
		metric_improvement = unoptimized_metric / optimized_metric if optimized_metric != 0 else 1

		# Calculate cycle latency improvement percentage
		cycle_latency_improvement = (
			(unoptimized_metric - optimized_metric) / unoptimized_metric * 100 if unoptimized_metric > 0 else 0
		)
		self.optimization_report = ""

		# Format the atomic contention improvement message
		if metric_improvement > 1:
			self.optimization_report += (
				f"Atomic Contention Reduction: Successfully reduced atomic contention by "
				f"{metric_improvement:.2f}x. "
				f"Average atomic latency improved from {unoptimized_metric:.0f} to {optimized_metric:.0f} cycles "
				f"({cycle_latency_improvement:.1f}% reduction - lower latency means less contention). "
			)
		else:
			self.optimization_report += (
				f"Atomic Contention Increase: Atomic contention increased by "
				f"{1 / metric_improvement:.2f}x. "
				f"Average atomic latency worsened from {unoptimized_metric:.0f} to {optimized_metric:.0f} cycles "
				f"({abs(cycle_latency_improvement):.1f}% increase - higher latency means more contention). "
			)

		# Format the performance improvement message
		if speedup > 1:
			self.optimization_report += (
				f"Performance Gain: Achieved {speedup:.2f}x speedup with execution time "
				f"reduced from {unoptimized_time / 1e6:.2f}ms to {optimized_time / 1e6:.2f}ms "
				f"({(speedup - 1) * 100:.1f}% faster)."
			)
		else:
			self.optimization_report += (
				f"Performance Loss: Experienced {1 / speedup:.2f}x slowdown with execution time "
				f"increased from {unoptimized_time / 1e6:.2f}ms to {optimized_time / 1e6:.2f}ms "
				f"({(1 / speedup - 1) * 100:.1f}% slower)."
			)

		if not success or speedup < 1:
			self.current_summary = self.optimization_report

			# Update the latest optimization attempt with results
			if self.optimization_attempts:
				latest_attempt = self.optimization_attempts[-1]
				latest_attempt["success"] = False
				latest_attempt["improvement"] = -cycle_latency_improvement  # Negative for worsening
				logging.info(
					f"Optimization iteration {latest_attempt['iteration']} FAILED: "
					f"Atomic contention increased by {abs(cycle_latency_improvement):.1f}%"
				)

			return Result(success=False, error_report=self.optimization_report)

		# Update the latest optimization attempt with results
		if self.optimization_attempts:
			latest_attempt = self.optimization_attempts[-1]
			latest_attempt["success"] = True
			latest_attempt["improvement"] = cycle_latency_improvement
			logging.info(
				f"Optimization iteration {latest_attempt['iteration']} SUCCEEDED: "
				f"Atomic contention reduced by {cycle_latency_improvement:.1f}%"
			)

		# Log performance validation results
		optimized_code_string = None
		if self.optimization_attempts:
			optimized_code_string = self.optimization_attempts[-1].get("code")
		self.get_logger().record(
			"performance_validation_complete",
			{
				"success": success,
				"unoptimized_time_ns": unoptimized_time,
				"optimized_time_ns": optimized_time,
				"unoptimized_metric": unoptimized_metric,
				"optimized_metric": optimized_metric,
				"speedup": speedup,
				"metric_improvement": metric_improvement,
				"cycle_latency_improvement": cycle_latency_improvement,
				"optimization_report": self.optimization_report,
				"optimized_code_string": optimized_code_string,
			},
		)

		logging.info(self.optimization_report)

		self.success = True
		return Result(success=True, asset={"log": self.optimization_report})

	def write_results(self, output_file: str = None):
		"""
		Writes the results to the output file.
		"""
		super().write_results(
			output_file=output_file,
			additional_results={"formula": "atomicContention", "success": self.success},
		)

	def summarize_previous_passes(self):
		"""
		Summarizes the results of the previous passes for future prompts.
		"""
		pass

	def get_optimization_history_context(self) -> str:
		"""
		Generate context about previous optimization attempts for the LLM.

		Returns:
		    str: Formatted context about previous optimization attempts
		"""
		if not self.optimization_attempts:
			return "This is the first optimization attempt."

		context = "PREVIOUS OPTIMIZATION ATTEMPTS:\n"
		for i, attempt in enumerate(self.optimization_attempts, 1):
			context += f"Iteration {i}: "
			if attempt["success"]:
				context += f"SUCCESS - Atomic contention reduced by {attempt['improvement']:.1f}%\n"
			else:
				context += f"FAILED - Atomic contention increased by {abs(attempt['improvement']):.1f}%\n"

		return context

	def reset_optimization_state(self):
		"""
		Reset the optimization state for a new optimization cycle.
		"""
		self.optimization_attempts = []
		self.iteration_count = 0
		self.current_summary = None
		self.previous_source_code = None
		self.success = False
		logging.info("Optimization state reset for new cycle")
