import json
import time
import demistomock as demisto
from datetime import datetime
from typing import Any
from CommonServerPython import *

THREATMON_PAGE_SIZE = 10
MAX_PAGES_PER_RUN = 100  # Max pages fetched per run (1000 incidents) to avoid timeout


class Client(BaseClient):
    def __init__(self, api_url: str, api_key: str, verify: bool, proxy: bool):
        super().__init__(
            base_url=api_url, verify=verify, proxy=proxy, headers={"X-COMPANY-API-KEY": api_key, "accept": "application/json"}
        )

    def get_incidents(self, last_incident_id=None, page=0):
        """Fetches incidents from Threatmon API using pagination and lastIncidentId filtering."""

        url_suffix = f"/vulnerabilities/{page}"
        params = {}
        if last_incident_id:
            params["afterAlarmCode"] = last_incident_id

        return self._http_request(method="GET", url_suffix=url_suffix, params=params)

    def set_status(self, data):
        """Updates incidents on Threatmon API using PATCH request."""

        return self._http_request(method="PATCH", url_suffix="/incident/status", json_data=data, resp_type="response")

    def request_takedown(self, finding_id: int, finding: str):
        """Submits a takedown request for a specific finding (alarm row)."""

        payload = {"findingId": finding_id, "finding": finding}
        return self._http_request(
            method="POST",
            url_suffix="/takedown",
            json_data=payload,
            ok_codes=(200, 400, 403, 404, 409),
            resp_type="response",
        )

    def request_data_removal(self, finding_id: int, finding: str):
        """Submits a Black Market Monitoring data removal request for a specific finding (alarm row)."""

        payload = {"findingId": finding_id, "finding": finding}
        return self._http_request(
            method="POST",
            url_suffix="/blackMarket/dataRemoval",
            json_data=payload,
            ok_codes=(200, 400, 403, 404, 409),
            resp_type="response",
        )

    def get_cve_list(self, page: int = 0, cvss: str | None = None):
        """Fetches a paginated list of all CVEs monitored by ThreatMon."""

        params = {}
        if cvss:
            params["cvss"] = cvss

        return self._http_request(method="GET", url_suffix=f"/cve/{page}", params=params)

    def get_subscribed_cve_list(self, page: int = 0, cvss: str | None = None, customer_name: str | None = None):
        """Fetches a paginated list of CVEs affecting products the company is subscribed to."""

        params = {}
        if cvss:
            params["cvss"] = cvss
        if customer_name:
            params["customerName"] = customer_name

        return self._http_request(method="GET", url_suffix=f"/cve/subscribed/{page}", params=params)


def convert_to_demisto_severity(severity: str) -> int:
    """Maps Threatmon severity to Cortex XSOAR/XSIAM severity (1 to 4)."""
    severity_mapping = {"Information": 1, "Low": 1, "Medium": 2, "High": 3, "Critical": 4}
    return severity_mapping.get(severity, 1)


def fetch_incidents(client: Client, last_run: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Fetches new incidents from Threatmon API using pagination and latest lastIncidentId logic.

    Fetches at most MAX_PAGES_PER_RUN pages per run to avoid timeout. Pagination state
    is saved in last_run so subsequent runs continue from where the previous run left off.
    last_incident_id is only advanced once the current batch is fully consumed to ensure
    consistent afterAlarmCode filtering across pages.
    """

    last_incident_id = last_run.get("last_incident_id")

    if last_incident_id is None:
        last_incident_id = demisto.params().get("lastIncidentId")

    if last_incident_id is None:
        last_incident_id = 0

    try:
        last_incident_id = int(last_incident_id)
    except (ValueError, TypeError):
        last_incident_id = 0

    try:
        page = int(last_run.get("page", 0))
    except (ValueError, TypeError):
        page = 0
    # max_seen_id tracks the highest alarmCode seen across pages within the current batch.
    # It is kept separate from last_incident_id so that afterAlarmCode stays constant
    # for all pages of a batch (changing it mid-batch would shift page offsets).
    try:
        max_seen_id = int(last_run.get("max_seen_id", last_incident_id))
    except (ValueError, TypeError):
        max_seen_id = last_incident_id
    incidents = []
    pages_fetched = 0

    while pages_fetched < MAX_PAGES_PER_RUN:
        response = client.get_incidents(last_incident_id=last_incident_id, page=page)

        alerts = response.get("data", [])

        if not alerts:
            # Batch complete: advance the filter cursor and reset page
            last_incident_id = max_seen_id
            page = 0
            last_run.pop("max_seen_id", None)
            break

        for alert in alerts:
            incident_id = 0
            raw_alarm_code = alert.get("alarmCode")
            if raw_alarm_code is not None:
                try:
                    incident_id = int(raw_alarm_code)
                except (ValueError, TypeError):
                    demisto.debug(f"Invalid alarmCode value received from API: {raw_alarm_code}. Defaulting incident_id to 0.")
            title = alert.get("title", "Unknown Threat")
            description = alert.get("description", "No description available")
            severity = alert.get("severity", "Low")
            status = alert.get("status", "New")
            alarm_date = alert.get("alarmDate", datetime.utcnow().isoformat())

            incident = {
                "name": f"Threatmon Alert: {title}",
                "details": description,
                "severity": convert_to_demisto_severity(severity),
                "occurred": alarm_date,
                "rawJSON": json.dumps(alert),
                "labels": [{"type": "Status", "value": status}],
            }
            incidents.append(incident)
            max_seen_id = max(max_seen_id, incident_id)

        pages_fetched += 1
        page += 1

        if len(alerts) < THREATMON_PAGE_SIZE:
            # Last page of batch: advance cursor and reset
            last_incident_id = max_seen_id
            page = 0
            last_run.pop("max_seen_id", None)
            break

        time.sleep(1)
    else:
        # Reached MAX_PAGES_PER_RUN with more pages remaining — save state for next run.
        # last_incident_id stays unchanged so afterAlarmCode filter remains consistent.
        last_run["max_seen_id"] = max_seen_id

    last_run["last_incident_id"] = last_incident_id
    last_run["page"] = page
    return incidents, last_run


def test_module(client: Client) -> str:
    """Tests API connectivity and authentication."""
    try:
        response = client.get_incidents(page=0)

        if isinstance(response, dict) and "data" in response:
            return "ok"

        if isinstance(response, dict) and "message" in response:
            return_error(f"Test failed: {response.get('message')}")

        return_error("Test failed: Unexpected API response structure.")

    except DemistoException as e:
        if "401" in str(e) or "403" in str(e):
            return_error("Test failed: Authentication error. Please check your API credentials.")
        elif "500" in str(e):
            return_error("Test failed: Server error (500). Please try again later.")
        else:
            return_error(f"Test failed: {str(e)}")

    except Exception as e:
        return_error(f"Test failed: {str(e)}")

    return "ok"


def change_incident_status(client: Client, args: dict[str, Any]) -> CommandResults:
    code = args.get("alarmId")
    status = args.get("status")
    data = {"status": status, "alarmIds": [code]}
    response = client.set_status(data=data)
    if response.ok:
        return CommandResults(readable_output=f"Incident {code} status changed to {status}.")
    else:
        raise DemistoException(
            f"Failed to change status for incident {code}. API returned {response.status_code}: {response.text}"
        )


def request_takedown_command(client: Client, args: dict[str, Any]) -> CommandResults:
    finding_id_raw = args.get("findingId")
    finding = args.get("finding")

    if not finding_id_raw:
        raise ValueError("findingId argument is required.")
    if not finding:
        raise ValueError("finding argument is required.")

    try:
        finding_id = int(finding_id_raw)
    except (ValueError, TypeError):
        raise ValueError(f"findingId must be a valid integer, got: {finding_id_raw}")

    response = client.request_takedown(finding_id=finding_id, finding=finding)

    status_code = response.status_code
    if status_code == 200:
        return CommandResults(readable_output="Takedown request submitted successfully.")
    elif status_code == 404:
        raise DemistoException(f"Finding not found: findingId={finding_id}")
    elif status_code == 409:
        raise DemistoException(f"A takedown request already exists for findingId={finding_id}")
    elif status_code == 403:
        raise DemistoException("Takedown quota exceeded. Please contact ThreatMon.")
    elif status_code == 400:
        raise DemistoException("This finding is not eligible for a takedown request.")
    else:
        raise DemistoException(f"API Error: {status_code} - {response.text}")


def request_data_removal_command(client: Client, args: dict[str, Any]) -> CommandResults:
    finding_id_raw = args.get("findingId")
    finding = args.get("finding")

    if not finding_id_raw:
        raise ValueError("findingId argument is required.")
    if not finding:
        raise ValueError("finding argument is required.")

    try:
        finding_id = int(finding_id_raw)
    except (ValueError, TypeError):
        raise ValueError(f"findingId must be a valid integer, got: {finding_id_raw}")

    response = client.request_data_removal(finding_id=finding_id, finding=finding)

    status_code = response.status_code
    if status_code == 200:
        return CommandResults(readable_output="Data removal request submitted successfully.")
    elif status_code == 404:
        raise DemistoException(f"Finding not found: findingId={finding_id}")
    elif status_code == 409:
        raise DemistoException(f"A data removal request already exists for findingId={finding_id}")
    elif status_code == 403:
        raise DemistoException("Data removal quota exceeded or insufficient rights. Please contact ThreatMon.")
    elif status_code == 400:
        raise DemistoException("This finding is not eligible for a data removal request.")
    else:
        raise DemistoException(f"API Error: {status_code} - {response.text}")


def format_cve_vendors(vendors: dict[str, list[str]]) -> str:
    """Formats a vendor-to-products map into a human-readable string for the markdown table."""

    if not vendors:
        return ""
    return "; ".join(f"{vendor}: {', '.join(products)}" for vendor, products in vendors.items())


def build_cve_command_results(response: dict[str, Any], readable_title: str) -> CommandResults:
    """Builds CommandResults for a page of CVE records returned by the ThreatMon CVE endpoints."""

    records = response.get("data") or []
    total_records = response.get("totalRecords", 0)

    markdown_rows = []
    for record in records:
        row = dict(record)
        row["vendors"] = format_cve_vendors(record.get("vendors") or {})
        markdown_rows.append(row)

    headers = [
        "cve",
        "summary",
        "cvssV2",
        "severityV2",
        "cvssV3",
        "severityV3",
        "cvssV3_1",
        "severityV3_1",
        "cvssV4",
        "severityV4",
        "vendors",
        "exploit",
        "knownRansomwareCampaignUse",
        "zeroday",
        "createdAt",
        "updatedAt",
    ]
    title = f"{readable_title} (Total Records: {total_records})"
    readable_output = tableToMarkdown(
        title,
        markdown_rows,
        headers=headers,
        headerTransform=string_to_table_header,
        removeNull=True,
    )

    return CommandResults(
        outputs_prefix="ThreatMon.CVE",
        outputs_key_field="cve",
        outputs=records,
        readable_output=readable_output,
        raw_response=response,
    )


def list_cves_command(client: Client, args: dict[str, Any]) -> CommandResults:
    page = arg_to_number(args.get("page")) or 0
    cvss = args.get("cvss")

    response = client.get_cve_list(page=page, cvss=cvss)
    return build_cve_command_results(response, "ThreatMon CVE List")


def list_subscribed_cves_command(client: Client, args: dict[str, Any]) -> CommandResults:
    page = arg_to_number(args.get("page")) or 0
    cvss = args.get("cvss")
    customer_name = args.get("customer_name")

    response = client.get_subscribed_cve_list(page=page, cvss=cvss, customer_name=customer_name)
    return build_cve_command_results(response, "ThreatMon Subscribed CVE List")


def main():
    """Main function called by Cortex XSOAR/XSIAM."""
    try:
        params = demisto.params()
        api_url = params.get("url", "https://external.threatmonit.io/api/threatmon/external/v1")
        credentials = params.get("credentials", {})
        api_key = credentials.get("password")
        verify = not params.get("insecure", False)
        proxy = params.get("proxy", False)

        client = Client(api_url=api_url, api_key=api_key, verify=verify, proxy=proxy)
        command = demisto.command()

        if command == "test-module":
            return_results(test_module(client))
        elif command == "threatmon_update_incident_status":
            return_results(change_incident_status(client, demisto.args()))
        elif command == "threatmon_request_takedown":
            return_results(request_takedown_command(client, demisto.args()))
        elif command == "threatmon_request_data_removal":
            return_results(request_data_removal_command(client, demisto.args()))
        elif command == "threatmon_list_cves":
            return_results(list_cves_command(client, demisto.args()))
        elif command == "threatmon_list_subscribed_cves":
            return_results(list_subscribed_cves_command(client, demisto.args()))
        elif command == "fetch-incidents":
            last_run = demisto.getLastRun() or {}
            incidents, last_run = fetch_incidents(client, last_run)
            demisto.setLastRun(last_run)
            demisto.incidents(incidents)

    except Exception as e:
        return_error(f"Error in Threatmon integration: {str(e)}")


if __name__ in ("__main__", "builtin", "builtins"):
    main()