from metalware_sdk.havoc_common_schema import *
from typing import List, Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    from metalware_sdk.havoc_client import HavocClient


class HavocAnalysis:
    """Represents an analysis instance.

    An analysis is the top-level container for runs, configuration,
    and firmware images.
    Hierarchy: Analysis -> [Run]
    """

    def __init__(self, client: 'HavocClient', analysis_id: str):
        self._client = client
        self._analysis_id = analysis_id

    @property
    def analysis_id(self) -> str:
        return self._analysis_id

    # === Status & Lifecycle ===

    def get_status(self) -> 'AnalysisStatusResponse':
        """Get the current status of this analysis."""
        return self._client.get_analysis_status(self._analysis_id)

    def stop(self) -> None:
        """Stop this analysis."""
        self._client.stop_analysis(self._analysis_id)

    def resume(self) -> None:
        """Resume this analysis (must be in Failed or Completed state)."""
        self._client.resume_analysis(self._analysis_id)

    def delete(self) -> None:
        """Delete this analysis and all its data."""
        self._client.delete_analysis(self._analysis_id)

    # === Run Management ===

    def start_run(self, workers: int = 1, dry_run: bool = False) -> str:
        """Start a new run for this analysis.

        Args:
            workers: Number of workers to use (default: 1)
            dry_run: Whether this is a dry run (default: False)

        Returns:
            The run_id (UUID string) of the newly started run.
        """
        run_params = RunParams(worker_count=workers, dry_run=dry_run)
        return self._client.start_run(self._analysis_id, run_params)

    def get_runs(self) -> List[str]:
        """Get all run IDs associated with this analysis.

        Returns:
            List of run_id (UUID strings) for runs in this analysis.
        """
        status = self.get_status()
        return status.run_ids

    def get_run_status(self, run_id: str) -> RunStatus:
        """Get status of a specific run.

        Args:
            run_id: The run UUID to check

        Returns:
            The RunStatus of the specified run.
        """
        return self._client.get_run_status(run_id)

    def stop_run(self, run_id: str) -> None:
        """Stop a specific run.

        Args:
            run_id: The run UUID to stop
        """
        self._client.stop_run(run_id)

    def get_run_stats(self, run_id: str) -> RunStats:
        """Get statistics for a specific run.

        Args:
            run_id: The run UUID to get stats for

        Returns:
            RunStats for the specified run.
        """
        return self._client.get_run_stats(run_id)

    def get_testcases(self, run_id: str) -> List[Testcase]:
        """Get testcases for a specific run.

        Args:
            run_id: The run UUID to get testcases for

        Returns:
            List of Testcase objects.
        """
        return self._client.get_testcases(run_id)

    # === Tags ===

    def set_tags(self, tags: list[str]) -> None:
        """Set tags on this analysis. Can be called in any stage."""
        self._client.set_tags(self._analysis_id, tags)

    def get_tags(self) -> list[str]:
        """Get tags for this analysis."""
        status = self.get_status()
        return getattr(status, 'tags', [])

    # === Fuzzer Configuration ===

    def set_fuzzer_config(self, fuzzer_config: FuzzerConfig) -> None:
        self._client.set_fuzzer_config(self._analysis_id, fuzzer_config)

    def get_fuzzer_config(self) -> FuzzerConfig:
        return self._client.get_fuzzer_config(self._analysis_id)

    # === Image Management ===

    def set_image(self, image: Image) -> None:
        self._client.set_image(self._analysis_id, image)

    def get_image(self) -> Image:
        return self._client.get_analysis_image(self._analysis_id)

    def add_raw_image(self, entry_address: int, segments: List[Tuple[int, str]]) -> Image:
        raw_image = RawImage(segments=[])
        for (address, path) in segments:
            file_metadata = self._client.upload_file(path)
            raw_image.segments.append(RawImageSegment(address=address, file_id=file_metadata.id, file_offset=0))
        image_config = Image(entry_address=entry_address, image_format=ImageFormat(Raw=raw_image))
        self._client.set_image(self._analysis_id, image_config)
        return image_config

    def add_elf_image(self, entry_address: int, file_path: str) -> Image:
        file_metadata = self._client.upload_file(file_path)
        image_config = Image(entry_address=entry_address, image_format=ImageFormat(Elf=file_metadata.id))
        self._client.set_image(self._analysis_id, image_config)
        return image_config

    # === Memory Map ===

    def get_memory_map(self) -> Dict[str, Memory]:
        mcu = self._client.get_mcu(self._analysis_id)
        if mcu.CortexMGeneric is not None:
            return mcu.CortexMGeneric.memory_map
        else:
            raise ValueError(f"Unsupported MCU type: {type(mcu)}")

    def set_memory_map(self, memory_map: Dict[str, Memory]) -> None:
        self._client.set_memory_map(self._analysis_id, memory_map)
