import httpx
import time
import hmac
import hashlib
from uuid import uuid4
from typing import Dict, Any, Optional, List
from types import TracebackType
from .exceptions import ArdentAPIError, ArdentAuthError, ArdentValidationError
import json


class ArdentClient:
    def __init__(
        self,
        public_key: str,
        secret_key: str,
        base_url: str = "https://ardentbackendwebappfinal.azurewebsites.net",
    ):
        super().__init__()
        if not public_key or not secret_key:
            raise ArdentValidationError("Both public and secret keys are required")

        self.public_key = public_key
        self.secret_key = secret_key
        self.base_url = base_url.rstrip("/")
        self.session_id = str(uuid4())
        self._client = httpx.Client(timeout=3000.0)

    def _sign_request(
        self,
        method: str,
        path: str,
        body: str = "",
        header_overrides: Optional[Dict[str, str]] = None,
    ) -> Dict[str, str]:
        timestamp = str(int(time.time()))
        message = f"{timestamp}{method}{path}{body}"

        signature = hmac.new(
            self.secret_key.encode("utf-8"), message.encode("utf-8"), hashlib.sha256
        ).hexdigest()

        headers = {
            "X-API-Key": self.public_key,
            "X-Signature": signature,
            "X-Timestamp": timestamp,
            "X-Session-ID": self.session_id,
            "Content-Type": "application/json",
        }

        # Merge header overrides if provided
        if header_overrides:
            headers.update(header_overrides)

        return headers

    def create_job(
        self, message: str, header_overrides: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        path = "/v1/jobs/createJob"
        body = {
            "userMessage": message,
        }

        try:
            json_body = json.dumps(body, separators=(",", ":"))

            headers = self._sign_request(
                method="POST",
                path=path,
                body=json_body,
                header_overrides=header_overrides,
            )

            response = self._client.post(f"{self.base_url}{path}", headers=headers, json=body)

            response.raise_for_status()

            if response.status_code == 201:  # Handle 201 Created specifically
                response_data = response.json()
                if not response_data:
                    raise ArdentAPIError("API returned empty response")

                # Ensure required fields are present
                required_fields = ["id", "files_share_name", "userID"]
                if not all(field in response_data for field in required_fields):
                    # Generate an ID if missing
                    if "id" not in response_data:
                        response_data["id"] = str(uuid4())
                    # Use empty string for missing share name
                    if "files_share_name" not in response_data:
                        response_data["files_share_name"] = ""
                    # Use provided userID if missing

                return response_data

            return response.json()

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )
        except json.JSONDecodeError as e:
            raise ArdentAPIError(f"Invalid JSON response from API: {str(e)}")

    def execute_job(
        self,
        jobID: str,
        message: str,
        files_share_name: str,
        userID: str,
        safe_mode: bool = False,
        header_overrides: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Execute a job with the given parameters."""
        path = "/v1/jobs/APIChat"  # Updated endpoint path
        body: Dict[str, Any] = {
            "jobID": jobID,
            "userMessage": message,
            "files_share_name": files_share_name,
            "userID": userID,
            "safeMode": safe_mode,
        }

        try:
            json_body = json.dumps(body, separators=(",", ":"))
            headers = self._sign_request(
                method="POST",
                path=path,
                body=json_body,
                header_overrides=header_overrides,
            )

            response = self._client.post(f"{self.base_url}{path}", headers=headers, json=body)
            response.raise_for_status()
            return response.json()

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            elif e.response.status_code == 402:
                raise ArdentAPIError("Out of credits")
            elif e.response.status_code == 403:
                raise ArdentAuthError("Missing required scope: job:execute")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )

    def create_and_execute_job(
        self,
        org_id: str,
        message: str,
        safe_mode: bool = False,
        header_overrides: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Create and execute a job in one operation."""
        # First create the job
        path = "/v1/jobs/createJob"
        create_body = {
            "userMessage": message,
            "org_id": org_id,
        }

        try:
            # Create job
            json_body = json.dumps(create_body, separators=(",", ":"))
            headers = self._sign_request(
                method="POST",
                path=path,
                body=json_body,
                header_overrides=header_overrides,
            )

            create_response = self._client.post(
                f"{self.base_url}{path}", headers=headers, json=create_body
            )
            create_response.raise_for_status()
            job = create_response.json()

            if not job:
                raise ArdentAPIError("Job creation failed - empty response")

            # Then execute the job
            execute_path = "/v1/jobs/APIChat"
            execute_body: Dict[str, Any] = {
                "jobID": job["id"],
                "userMessage": message,
                "files_share_name": job["files_share_name"],
                "safeMode": safe_mode,
            }

            json_body = json.dumps(execute_body, separators=(",", ":"))
            headers = self._sign_request(
                method="POST",
                path=execute_path,
                body=json_body,
                header_overrides=header_overrides,
            )

            execute_response = self._client.post(
                f"{self.base_url}{execute_path}", headers=headers, json=execute_body
            )
            execute_response.raise_for_status()
            execute_result = execute_response.json()
            create_and_execute_responses = {**execute_result, **job}
            return create_and_execute_responses

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            elif e.response.status_code == 402:
                raise ArdentAPIError("Out of credits")
            elif e.response.status_code == 403:
                raise ArdentAuthError("Missing required scope: job:execute")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )

    def recursive_structure_validation(
        self, value: Any, structure_def: Any, path: str = ""
    ) -> None:
        """
        Recursively validate data against a structure definition.

        Args:
            value: The value to validate
            structure_def: The structure definition to validate against
            path: The path to the value in the structure
        """

        if isinstance(structure_def, type):
            if not isinstance(value, structure_def):
                raise ArdentValidationError(
                    f"Invalid type at {path}. Expected {structure_def.__name__}, got {type(value).__name__}"
                )
        elif isinstance(structure_def, dict):
            expected_type = structure_def["type"]
            if not isinstance(value, expected_type):
                raise ArdentValidationError(
                    f"Invalid type at {path}. Expected {expected_type.__name__}, got {type(value).__name__}"
                )

            if "structure" in structure_def:
                for i, item in enumerate(value):
                    item_path = f"{path}[{i}]"

                    for req_field in structure_def["structure"]["required"]:
                        if req_field not in item:
                            raise ArdentValidationError(
                                f"Missing required field '{req_field}' at {item_path}"
                            )

                    for prop_name, prop_def in structure_def["structure"]["properties"].items():
                        if prop_name in item:
                            self.recursive_structure_validation(
                                item[prop_name], prop_def, f"{item_path}.{prop_name}"
                            )

    def _validate_connector_structure(
        self,
        connector_type: str,
        connector_info: Dict[str, Any],
        connector_data: Dict[str, Any],
    ) -> None:
        """
        Validate connector structure based on connector type.
        """
        # First check required parameters
        required_params = connector_info["required_params"]
        for param in required_params:
            if param not in connector_data:
                raise ArdentValidationError(f"Missing required parameter: {param}")

        # Validate each field against its structure definition
        for param, structure_def in connector_info["structure"].items():
            if param in connector_data:
                self.recursive_structure_validation(connector_data[param], structure_def, param)

    def setup_connector(
        self,
        service_name: str,
        connection_details: Dict[str, Any],
        name: Optional[str] = None,
        selected_paths: Optional[List[str]] = None,
        header_overrides: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """
        Setup a complete data connector using v2 API.
        Complete flow: create -> discover -> select -> index

        Args:
            service_name: The service type ('snowflake', 'postgresql', 'mongodb', etc.)
            connection_details: Dictionary containing all connection parameters
            name: Optional friendly name for the connector
            selected_paths: Optional list of paths to select for indexing
            header_overrides: Optional dictionary of headers to override or add to the request

        Returns:
            Dict containing the full connector object with id for future operations
        """

        try:
            # Step 1: CREATE CONNECTOR - POST /v1/connectors
            create_body: Dict[str, Any] = {
                "service_name": service_name,
                "connection_details": connection_details,
            }
            if name:
                create_body["name"] = name

            create_path = "/v1/connectors"
            json_body = json.dumps(create_body, separators=(",", ":"))
            headers = self._sign_request(
                method="POST",
                path=create_path,
                body=json_body,
                header_overrides=header_overrides,
            )

            response = self._client.post(
                f"{self.base_url}{create_path}", headers=headers, json=create_body
            )
            response.raise_for_status()
            connector = response.json()
            connector_id = connector["id"]  # Backend generates this

            # Step 2: DISCOVER CONNECTOR - POST /v1/connectors/{connector_id}/discover
            discover_path = f"/v1/connectors/{connector_id}/discover"
            discover_headers = self._sign_request(
                method="POST",
                path=discover_path,
                body="{}",
                header_overrides=header_overrides,
            )

            discover_response = self._client.post(
                f"{self.base_url}{discover_path}", headers=discover_headers, json={}
            )
            discover_response.raise_for_status()

            # Step 3: SELECT PATHS (if provided) - POST /v1/connectors/{connector_id}/selection
            if selected_paths:
                select_body = {"selected_paths": selected_paths}
                select_path = f"/v1/connectors/{connector_id}/selection"
                select_json_body = json.dumps(select_body, separators=(",", ":"))
                select_headers = self._sign_request(
                    method="POST",
                    path=select_path,
                    body=select_json_body,
                    header_overrides=header_overrides,
                )

                select_response = self._client.post(
                    f"{self.base_url}{select_path}",
                    headers=select_headers,
                    json=select_body,
                )
                select_response.raise_for_status()

            # Step 4: INDEX CONNECTOR - POST /v1/connectors/{connector_id}/index
            index_path = f"/v1/connectors/{connector_id}/index"
            index_headers = self._sign_request(
                method="POST",
                path=index_path,
                body="{}",
                header_overrides=header_overrides,
            )

            index_response = self._client.post(
                f"{self.base_url}{index_path}", headers=index_headers, json={}
            )
            index_response.raise_for_status()

            return connector  # Return the full connector object with id

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )

    def delete_connector(
        self, connector_id: str, header_overrides: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """
        Delete a connector by its ID using v1 API.

        Args:
            connector_id: The ID of the connector to delete
            header_overrides: Optional dictionary of headers to override or add to the request
        """
        path = f"/v1/connectors/{connector_id}"

        try:
            headers = self._sign_request(
                method="DELETE",
                path=path,
                body="",  # Empty body for DELETE
                header_overrides=header_overrides,
            )

            response = self._client.request("DELETE", f"{self.base_url}{path}", headers=headers)
            response.raise_for_status()

            # DELETE operations typically return 204 No Content or success message
            if response.status_code == 204:
                return {
                    "status": "success",
                    "message": "Connector deleted successfully",
                }

            try:
                return response.json()
            except json.JSONDecodeError:
                return {
                    "status": "success",
                    "message": "Connector deleted successfully",
                }

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            elif e.response.status_code == 403:
                raise ArdentAuthError("Missing required scope: connectors.delete")
            elif e.response.status_code == 404:
                raise ArdentAPIError("Connector not found")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )

    def delete_job(
        self, job_id: str, header_overrides: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """
        Delete a job by its ID.

        Args:
            job_id: The ID of the job to delete
            header_overrides: Optional dictionary of headers to override or add to the request

        Returns:
            Dict containing confirmation of deletion
        """
        path = "/v1/jobs/deleteJob"
        body = {
            "jobID": job_id,
        }

        try:
            json_body = json.dumps(body, separators=(",", ":"))
            headers = self._sign_request(
                method="DELETE",
                path=path,
                body=json_body,
                header_overrides=header_overrides,
            )

            response = self._client.request(
                "DELETE", f"{self.base_url}{path}", headers=headers, json=body
            )
            response.raise_for_status()

            # DELETE operations typically return 204 No Content
            if response.status_code == 204:
                return {"status": "success", "message": "Job deleted successfully"}

            # Handle any other success responses
            try:
                return response.json()
            except json.JSONDecodeError:
                return {"status": "success", "message": "Job deleted successfully"}

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            elif e.response.status_code == 403:
                raise ArdentAuthError(
                    "Missing required scope: job:delete or cannot delete jobs for different user"
                )
            elif e.response.status_code == 404:
                raise ArdentAPIError("Job not found")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )

    def close(self):
        """Close the underlying HTTP client and clean up resources."""
        if hasattr(self, "_client"):
            self._client.close()

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> bool:
        """Context manager exit - ensures cleanup."""
        self.close()
        return False

    def get_job_files(
        self, jobID: str, header_overrides: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """
        Get the file system associated with a specific job.

        Args:
            jobID: The ID of the job to get files for
            header_overrides: Optional dictionary of headers to override or add to the request

        Returns:
            Dict containing the file system information
        """
        path = f"/v1/jobs/{jobID}/files"

        try:
            headers = self._sign_request(method="GET", path=path, header_overrides=header_overrides)

            response = self._client.get(f"{self.base_url}{path}", headers=headers)
            response.raise_for_status()
            return response.json()

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ArdentAuthError("Invalid API key or signature")
            elif e.response.status_code == 404:
                raise ArdentAPIError("Job not found")
            raise ArdentAPIError(
                f"API request failed: {str(e)}",
                status_code=e.response.status_code,
                response=e.response,
            )
