import httpx
import time
import hmac
import hashlib
from uuid import uuid4
from typing import Dict, Any, Optional, List
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",
    ):
        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, org_id: str, header_overrides: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """
        Create a new job.
        
        Args:
            message: The user message/instruction for the job
            org_id: Organization ID (required by backend)
            header_overrides: Optional dictionary of headers to override or add to the request
            
        Returns:
            Dict containing job creation response with id, files_share_name, userID
        """
        path = "/v1/jobs/createJob"
        body = {
            "userMessage": message,
            "org_id": org_id,
        }

        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 = {
            "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,
        message: str,
        org_id: str,
        safe_mode: bool = False,
        header_overrides: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """
        Create and execute a job in one operation.
        
        Args:
            message: The user message/instruction for the job
            org_id: Organization ID (required by backend)
            safe_mode: Whether to run in safe mode (default: False)
            header_overrides: Optional dictionary of headers to override or add to the request
            
        Returns:
            Dict containing combined job creation and execution response
        """
        # 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 = {
                "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 /v2/connectors
            create_body = {
                "service_name": service_name,
                "connection_details": connection_details,
            }
            if name:
                create_body["name"] = name

            create_path = "/v2/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 /v2/connectors/{connector_id}/discover
            discover_path = f"/v2/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 /v2/connectors/{connector_id}/selection
            if selected_paths:
                select_body = {"selected_paths": selected_paths}
                select_path = f"/v2/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 /v2/connectors/{connector_id}/index
            index_path = f"/v2/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 v2 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"/v2/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, exc_val, exc_tb):
        """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,
            )
