from metalware_sdk.havoc_common_schema import *
from metalware_sdk.havoc_extras import TestcaseInput
import requests
import base64
from typing import Optional, Tuple, List, Dict, Union, TYPE_CHECKING
from dataclasses import dataclass, field
import os

if TYPE_CHECKING:
    from metalware_sdk.havoc_analysis import HavocAnalysis


@dataclass
class HavocClient:
  """Client for interacting with the Havoc web server API."""
  base_url: str
  verify_ssl: bool = True
  token: Optional[str] = None
  session: requests.Session = field(default_factory=requests.Session)

  def __post_init__(self):
    if self.token:
      self.session.headers["Authorization"] = f"Bearer {self.token}"

  def _make_request(self, method: str, endpoint: str, **kwargs) -> requests.Response:
    url = f"{self.base_url}/api/{endpoint.lstrip('/')}"
    kwargs.setdefault('verify', self.verify_ssl)
    try:
      resp = self.session.request(method, url, **kwargs)
      resp.raise_for_status()
      return resp
    except requests.exceptions.RequestException as e:
      raise RuntimeError(f"Request to {url} failed: {str(e)}.")

  def get_analyses(self) -> list:
    resp = self._make_request('GET', '/analyses')
    return resp.json()

  def upload_file(self, file_path: str, label: str = "unnamed") -> FileMetadata:
    if not os.path.exists(file_path):
      raise FileNotFoundError(f"File not found: {file_path}")

    with open(file_path, 'rb') as f:
      file_data = f.read()

    # Encode file data as base64
    encoded_data = base64.b64encode(file_data)

    # Make request
    resp = self._make_request(
      'POST',
      '/upload-file',
      params={'label': label},
      data=encoded_data
    )

    result = resp.json()
    if isinstance(result, dict) and 'Ok' in result:
      return FileMetadata(**result['Ok'])
    else:
      raise RuntimeError(f"Upload failed: {result.get('Err', 'Unknown error')}")

  def infer_config(self, file_id: str, mcu: str) -> InferredConfig:
    resp = self._make_request(
      'POST',
      f'/infer-config',
      json=InferConfigRequest(file_id=file_id, mcu=mcu).to_dict()
    )

    result = resp.json()
    if isinstance(result, dict) and 'Ok' in result:
      ic = InferredConfig.from_dict(result['Ok'])
      return ic
    else: raise RuntimeError(f"Memory config inference failed: {result.get('Err', 'Unknown error')}")

  def infer_config_from_file(self, file_path: str, mcu: str) -> InferredConfig:
    # upload first
    file_metadata = self.upload_file(file_path)
    return self.infer_config(file_id=file_metadata.id, mcu=mcu)

  def create_analysis(self, mcu, image: Image, auto: bool = True, worker_count: int | None = None, timeout_secs: int | None = None, tags: list[str] | None = None) -> 'HavocAnalysis':
    """Create a new analysis. Returns a HavocAnalysis wrapper.

    Args:
        mcu: MCU configuration (CortexMGeneric, TI_CC430F6137, or Mcu instance).
        image: Image configuration.
        auto: If True (default), the background state machine drives execution.
              If False, users control runs manually via API.
        worker_count: Number of fuzzer workers for this analysis.
                      If None, uses the server's global workers_per_analysis setting.
        timeout_secs: Analysis timeout in seconds. 0 disables the timeout.
                      If None, uses the server's global analysis_timeout_secs setting.
        tags: Optional list of tags for organizing/grouping this analysis.

    Returns:
        A HavocAnalysis wrapper around the new analysis.
    """
    if isinstance(mcu, CortexMGeneric):
      mcu = Mcu(CortexMGeneric=mcu)
    elif isinstance(mcu, TI_CC430F6137):
      mcu = Mcu(TI_CC430F6137=mcu)
    elif not isinstance(mcu, Mcu):
      raise ValueError(f"Unsupported MCU type: {type(mcu)}")

    if image is None:
      raise ValueError("image is required for create_analysis")

    body = {"mcu": mcu.to_dict(), "image": image.to_dict(), "auto": auto}
    if worker_count is not None:
      body["worker_count"] = worker_count
    if timeout_secs is not None:
      body["timeout_secs"] = timeout_secs
    if tags is not None:
      body["tags"] = tags
    resp = self._make_request('POST', '/create-analysis', json=body)

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Analysis creation failed: {result['Err']}")
    else:
      analysis_id = result['Ok']['analysis_id']
      from metalware_sdk.havoc_analysis import HavocAnalysis
      return HavocAnalysis(self, analysis_id)

  def analysis_exists(self, analysis_id: str) -> bool:
    try:
      resp = self._make_request('GET', f'/analysis/{analysis_id}/status')
      result = resp.json()
      if isinstance(result, dict) and 'Err' in result:
        return False
      return True
    except RuntimeError:
      return False

  def analysis_image_exists(self, analysis_id: str) -> bool:
    try:
      self._make_request('GET', f'/analysis/{analysis_id}/image')
      return True
    except RuntimeError:
      return False

  def delete_analysis(self, analysis_id: str) -> None:
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/delete'
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Analysis deletion failed: {result['Err']}")

  def set_tags(self, analysis_id: str, tags: list[str]) -> None:
    """Set tags on an analysis. Tags can be edited in any stage."""
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/tags',
      json={"tags": tags}
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Failed to set tags: {result['Err']}")

  def set_fuzzer_config(self, analysis_id: str, fuzzer_config: FuzzerConfig) -> None:
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/fuzzer-config',
      json=fuzzer_config.to_dict()
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Fuzzer config setting failed: {result['Err']}")
    else: return result['Ok']

  def get_fuzzer_config(self, analysis_id: str) -> FuzzerConfig:
    resp = self._make_request(
      'GET',
      f'/analysis/{analysis_id}/fuzzer-config'
    )
    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Fuzzer config retrieval failed: {result['Err']}")
    else: return FuzzerConfig.from_dict(result['Ok'])

  def set_image(self, analysis_id: str, image: Image) -> None:
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/image',
      json=image.to_dict()
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Image setting failed: {result['Err']}")
    else: return result['Ok']

  def get_analysis_image(self, analysis_id: str) -> Image:
    resp = self._make_request(
      'GET',
      f'/analysis/{analysis_id}/image'
    )
    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Image retrieval failed: {result['Err']}")
    else: return Image.from_dict(result['Ok'])

  def get_analysis_status(self, analysis_id: str) -> AnalysisStatusResponse:
    """Get the status of an analysis by analysis_id (UUID)."""
    resp = self._make_request(
      'GET',
      f'/analysis/{analysis_id}/status'
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Analysis status retrieval failed: {result['Err']}")
    else: return AnalysisStatusResponse.from_dict(result['Ok'])

  def stop_analysis(self, analysis_id: str) -> None:
    """Stop an analysis by analysis_id (UUID)."""
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/stop'
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Analysis stop failed: {result['Err']}")

  def resume_analysis(self, analysis_id: str) -> None:
    """Resume a completed or failed analysis by analysis_id (UUID)."""
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/resume'
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Analysis resume failed: {result['Err']}")

  def get_analysis(self, analysis_id: str) -> 'HavocAnalysis':
    """Get an analysis object by analysis_id (UUID)."""
    from metalware_sdk.havoc_analysis import HavocAnalysis
    # Verify the analysis exists by getting its status
    self.get_analysis_status(analysis_id)
    return HavocAnalysis(self, analysis_id)

  def start_run(self, analysis_id: str, params: RunParams) -> str:
    """Start a run for an analysis. Returns the run_id (UUID string).

    Args:
        analysis_id: The analysis UUID to start the run against
        params: RunParams with worker_count and dry_run settings
    """
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/start-run',
      json=params.to_dict()
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Run start failed: {result['Err']}")
    else: return result['Ok']

  def get_run_status(self, run_id: str) -> RunStatus:
    """Get status of a run by run_id (UUID)."""
    resp = self._make_request(
      'GET',
      f'/run/{run_id}/summary'
    )
    result = resp.json()
    return RunSummary.from_dict(result).status

  def stop_run(self, run_id: str) -> None:
    """Stop a run by run_id (UUID)."""
    resp = self._make_request(
      'POST',
      f'/run/{run_id}/stop'
    )
    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Stop run failed: {result['Err']}")

  def get_runs(self, analysis_id: str) -> List[Tuple[str, RunSummary]]:
    """Get all runs for an analysis. Returns list of (run_id, RunSummary) tuples."""
    resp = self._make_request(
      'GET',
      f'/analysis/{analysis_id}/runs'
    )
    return [(run_id, RunSummary.from_dict(run)) for (run_id, run) in resp.json()]

  def get_run_stats(self, run_id: str) -> RunStats:
    """Get stats for a run by run_id (UUID)."""
    resp = self._make_request(
      'GET',
      f'/run/{run_id}/stats'
    )
    return RunStats.from_dict(resp.json())

  def get_testcases(self, run_id: str) -> List[Testcase]:
    """Get testcases for a run by run_id (UUID)."""
    resp = self._make_request(
      'GET',
      f'/run/{run_id}/testcases'
    )
    return [Testcase.from_dict(testcase) for testcase in resp.json()]

  def get_testcase_input(self, run_id: str, testcase_id: str) -> TestcaseInput:
    """Get testcase input by run_id (UUID) and testcase_id."""
    resp = self._make_request(
      'GET',
      f'/run/{run_id}/testcase/{testcase_id}/input'
    )
    return TestcaseInput.from_bytes(resp.content)

  def start_debug_session(self, run_id: str, testcase_id: str) -> None:
    """Start a debug session by run_id (UUID) and testcase_id."""
    resp = self._make_request(
      'POST',
      f'/run/{run_id}/debug-session/{testcase_id}/start'
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Debug session start failed: {result['Err']}")

  def send_debug_command(self, run_id: str, testcase_id: str, command: str) -> None:
    """Send a debug command by run_id (UUID) and testcase_id."""
    resp = self._make_request(
      'POST',
      f'/run/{run_id}/debug-session/{testcase_id}/command',
      json=command
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Debug command send failed: {result['Err']}")
    else: return result['Ok']

  def inject_image(self, zip_path: str) -> None:
    if not os.path.exists(zip_path):
      raise FileNotFoundError(f"File not found: {zip_path}")

    with open(zip_path, 'rb') as f:
      zip_data = f.read()

    resp = self._make_request(
      'POST',
      f'/inject-image',
      data=zip_data,
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Image injection failed: {result['Err']}")

  def inject_analysis(self, zip_path: str) -> str:
    """Inject an analysis from a zip file. The zip must contain a single
    analysis_<uuid>/ directory with a state.json and run directories.

    Returns:
        The analysis_id (UUID string) of the injected analysis.
    """
    if not os.path.exists(zip_path):
      raise FileNotFoundError(f"File not found: {zip_path}")

    with open(zip_path, 'rb') as f:
      zip_data = f.read()

    resp = self._make_request(
      'POST',
      '/inject-analysis',
      data=zip_data,
    )

    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Analysis injection failed: {result['Err']}")
    return result['Ok']['analysis_id']

  def get_mcu(self, analysis_id: str) -> Mcu:
    resp = self._make_request(
      'GET',
      f'/analysis/{analysis_id}/status'
    )
    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"MCU config retrieval failed: {result['Err']}")
    else: return Mcu.from_dict(result['Ok']['mcu'])

  def get_dashboard_status(self) -> dict:
    """Get dashboard status (system metrics, worker pool, per-analysis summaries, queue)."""
    resp = self._make_request('GET', '/dashboard')
    return resp.json()

  def update_settings(self, **kwargs) -> dict:
    """Update dashboard settings (worker_pool_size, workers_per_analysis, analysis_timeout_secs)."""
    resp = self._make_request('POST', '/dashboard/settings', json=kwargs)
    return resp.json()

  def set_memory_map(self, analysis_id: str, memory_map: Dict[str, Memory]) -> None:
    # Convert Memory objects to dictionaries for JSON serialization
    memory_map_dict = {k: v.to_dict() for k, v in memory_map.items()}
    resp = self._make_request(
      'POST',
      f'/analysis/{analysis_id}/memory-map',
      json=memory_map_dict
    )
    result = resp.json()
    if isinstance(result, dict) and 'Err' in result:
      raise RuntimeError(f"Memory map setting failed: {result['Err']}")
