from pathlib import Path
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from urllib.parse import urlparse
import os
import enum
import json
import time
import threading
from pydantic import BaseModel
from typing import Optional, List
from blocks_control_sdk.blocks_events.context import BLOCKS_EVENT__PROVIDER
from blocks_control_sdk.utils import BlocksRuntimeConfigKeys, get_blocks_runtime_config, patch_blocks_runtime_config, update_git_credential


BLOCKS_AUTH__JOB_SESSION_TOKEN = os.getenv("BLOCKS_AUTH__JOB_SESSION_TOKEN")
BLOCKS_CLIENTS__ORCHESTRATOR_URL = os.getenv("BLOCKS_CLIENTS__ORCHESTRATOR_URL")
BLOCKS_CHAT_ID = os.getenv("BLOCKS_CHAT_ID")
BLOCKS_TASK_ID = os.getenv("BLOCKS_TASK_ID")

print(f"BLOCKS_AUTH__JOB_SESSION_TOKEN: {BLOCKS_AUTH__JOB_SESSION_TOKEN}")

class BlocksUserTelemetryMessage(enum.Enum):
    USER_ACTIVITY = "USER_ACTIVITY"
    WAKE_UP = "WAKE_UP"
    CODE_EDITOR_EXPANDED = "CODE_EDITOR_EXPANDED"
    CODE_EDITOR_OPENED = "CODE_EDITOR_OPENED"

class BlocksMessageCreate(BaseModel):
    message: str
    role: str
    type: str
    chat_thread_id: Optional[str] = None
    artifact_id: Optional[str] = None
    artifact_metadata: Optional[dict] = None

class BlocksMessage(BaseModel):
    id: str
    role: str
    type: str
    task_id: str
    message: str
    attachment_file_name: Optional[str]
    attachment_file_location: Optional[str]
    artifact_id: Optional[str] = None
    artifact_metadata: Optional[dict] = None
    chat_thread_id: Optional[str] = None
    ts: float
    is_message_snapshot_resume: bool = False
    source_provider: Optional[str] = BLOCKS_EVENT__PROVIDER.WEBHOOK.value

class BlocksOAuthCredentialsValues(BaseModel):
    access_token: str
    refresh_token: Optional[str] = None
    expires_at_in_seconds: Optional[float] = None


class BaseAPI:
    def __init__(self, timeout=5):
        self.last_processed_ts = 0

        # Global timeout can be overridden by environment variable
        self.timeout = float(os.getenv("BLOCKS_HTTP_TIMEOUT", timeout))

        # Configure retry strategy
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[500, 502, 503, 504],
            allowed_methods=["GET", "POST"]
        )

        # Create session with retry adapter
        self.session = requests.Session()
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

        # Set global timeout for all requests
        # This creates a default that can be overridden per-request
        self.session.request = self._wrap_request_with_timeout(self.session.request)

    def _wrap_request_with_timeout(self, original_request):
        def request_with_timeout(*args, **kwargs):
            # Only set timeout if not explicitly provided
            if 'timeout' not in kwargs:
                kwargs['timeout'] = self.timeout
            return original_request(*args, **kwargs)
        return request_with_timeout

    def re_create_session(self):
        self.session = requests.Session()
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[500, 502, 503, 504],
            allowed_methods=["GET", "POST"]
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

        # Reapply global timeout wrapper
        self.session.request = self._wrap_request_with_timeout(self.session.request)
    
    def get_base_url(self, url: str) -> str:
        """
        Returns the base url from a given url

        Example:
        https://api.dev.blocksorg.com/v1/chats/123/messages -> https://api.dev.blocksorg.com
        """
        parsed = urlparse(url)
        return f"{parsed.scheme}://{parsed.netloc}"

class Orchestrator(BaseAPI):

    def __init__(self):
        super().__init__()
        self.base_url = self.get_base_url(BLOCKS_CLIENTS__ORCHESTRATOR_URL)

    def update_task(self, body: dict = {}):
        res = self.session.put(
            f"{self.base_url}/job/v1/tasks/{BLOCKS_TASK_ID}", 
            json=body, 
            headers={"Authorization": f"Bearer {BLOCKS_AUTH__JOB_SESSION_TOKEN}"
        })
        res.raise_for_status()
        return res.json()

    def send_idle_status(self, wait_until: int = 1) -> None:
        """Send idle status in a fire-and-forget manner to handle VM snapshotting gracefully."""
        def _send_status():
            try:
                print(f"Waiting for {wait_until} seconds before sending idle status...")
                time.sleep(wait_until)
                res = requests.put(
                    f"{self.base_url}/job/v1/tasks/{BLOCKS_TASK_ID}",
                    json={"status": "IDLE"},
                    headers={"Authorization": f"Bearer {BLOCKS_AUTH__JOB_SESSION_TOKEN}"},
                    timeout=1  # Timeout to prevent hanging during snapshot
                )
                res.raise_for_status()
                print("Successfully sent idle status")
            except requests.Timeout:
                print("send_idle_status timed out (VM likely snapshotted)")
            except Exception as e:
                print(f"send_idle_status error (expected): {e}")

        # Start daemon thread and return immediately
        thread = threading.Thread(target=_send_status, daemon=True)
        thread.start()
        return thread

    def get_upload_snapshot_url(self) -> str:
        res = self.session.put(
                f"{self.base_url}/job/v1/tasks/snapshot/upload",
                headers={
                    "Authorization": f"Bearer {BLOCKS_AUTH__JOB_SESSION_TOKEN}",
                }
            )
        try:
            res.raise_for_status()
        except Exception as e:
            text = res.text
            print(f"===> ERROR GETTING UPLOAD SNAPSHOT URL <===\n\n{text}\n\n=====END OF ERROR GETTING UPLOAD SNAPSHOT URL=====")
            raise
        return res.json().get("url")

    def upload_to_snapshot_url(self, url: str, abs_file_path: str) -> None:
        with open(abs_file_path, "rb") as f:
            res = self.session.put(url, data=f)
            res.raise_for_status()
        self.update_task({
            "is_snapshotted": True 
        })

    def download_from_presigned_url(self, presigned_url: str, abs_file_path: str) -> None:
        """Download content from a presigned URL to a specified file path"""
        res = self.session.get(presigned_url, stream=True)
        res.raise_for_status()
        
        # Ensure directory exists
        Path(abs_file_path).parent.mkdir(parents=True, exist_ok=True)
        
        # Download in chunks
        with open(abs_file_path, 'wb') as f:
            for chunk in res.iter_content(chunk_size=8192):
                f.write(chunk)


class Chat(BaseAPI):
    oauth_credentials_mapping: dict[str, BlocksOAuthCredentialsValues] = {}

    def __init__(self):
        super().__init__()
        self.last_processed_ts = 0
        self.base_url = self.get_base_url(BLOCKS_CLIENTS__ORCHESTRATOR_URL)

    def update_chat(self, body: dict = {}):
        res = self.session.put(f"{self.base_url}/job/v1/chats/{BLOCKS_CHAT_ID}", 
        json=body, 
        headers={
            "Authorization": f"Bearer {BLOCKS_AUTH__JOB_SESSION_TOKEN}"
        })
        res.raise_for_status()
        return res.json()

    def create_message(self, blocks_message: BlocksMessageCreate) -> BlocksMessage:

        body = {
            "message": blocks_message.message,
            "role": blocks_message.role,
            "type": blocks_message.type,
            "chat_thread_id": blocks_message.chat_thread_id
        }

        if blocks_message.artifact_id:
            body["artifact_id"] = blocks_message.artifact_id
        if blocks_message.artifact_metadata:
            body["artifact_metadata"] = blocks_message.artifact_metadata

        if not blocks_message.chat_thread_id:
            blocks_message.chat_thread_id = get_blocks_runtime_config().get(BlocksRuntimeConfigKeys.CHAT_THREAD_ID)
        
        if not blocks_message.chat_thread_id:
            BLOCKS_INITIAL_CHAT_THREAD_ID = os.getenv("BLOCKS_INITIAL_CHAT_THREAD_ID")
            blocks_message.chat_thread_id = BLOCKS_INITIAL_CHAT_THREAD_ID

        res = self.session.post(f"{self.base_url}/job/v1/chats/{BLOCKS_CHAT_ID}/messages", json=body, headers={
            "Authorization": f"Bearer {BLOCKS_AUTH__JOB_SESSION_TOKEN}"
        })
        res.raise_for_status()
        return BlocksMessage(**res.json())

    def send_code_server_password(self, get_chat_thread_id: callable) -> None:
        def _send_password():
            try:
                while True:
                    code_server_token = None
                    code_server_token_file = Path("/home/user/.config/code-server/.mytoken")
                    if not code_server_token_file.exists():
                        time.sleep(1)
                    else:
                        code_server_token = code_server_token_file.read_text().strip()
                        break
                    
                chat_thread_id = get_chat_thread_id()
                
                self.create_message(BlocksMessageCreate(
                    message=json.dumps({
                        "code_server_password": code_server_token
                    }),
                    role="assistant",
                    type="telemetry",
                    chat_thread_id=chat_thread_id
                ))
                
            except requests.Timeout:
                print("CODESERVER password send timed out")
            except Exception as e:
                print(f"CODESERVER password send error: {e}")

        # Start daemon thread and return immediately
        thread = threading.Thread(target=_send_password, daemon=True)
        thread.start()
        return thread

    def send_ssh_key_as_message(self, chat_thread_id: str):
        """Send SSH key as a message in a fire-and-forget manner."""
        def _send_ssh_key():
            try:
                # get contents from /home/user/task.pem
                ssh_key = Path("/home/user/task.pem").read_text()
                self.create_message(BlocksMessageCreate(
                    message=ssh_key,
                    role="assistant",
                    type="ssh_key",
                    chat_thread_id=chat_thread_id
                ))
                print("Successfully sent SSH key as message")
            except Exception as e:
                print(f"send_ssh_key_as_message error: {e}")

        # Start daemon thread and return immediately
        thread = threading.Thread(target=_send_ssh_key, daemon=True)
        thread.start()

    def get_messages(self) -> List[BlocksMessage]:
        """
        Messages are paginated, as:

        {
            "oauth_credentials": [
            {
                    "provider": "claude",
                    "credentials": {
                        "access_token": "sk-ant...",
                        "refresh_token": "sk-ant...",
                        "expires_at_in_seconds": 1157081633 
                    }
                },
            ],
            "items": [...],
                "total": 1,
                "page": 1,
                "total_pages": 100
            }
        }
        """
        all_messages = []
        page = 1
        highest_ts = self.last_processed_ts

        if not BLOCKS_CHAT_ID:
            return []
        
        while True:
            res = self.session.get(f"{self.base_url}/job/v1/chats/{BLOCKS_CHAT_ID}/messages", params={
                "gts": self.last_processed_ts,
                "page": page,
                "limit": 50
            }, headers={
                "Authorization": f"Bearer {BLOCKS_AUTH__JOB_SESSION_TOKEN}"
            })
            res.raise_for_status()
            data = res.json()

            # Update the oauth credentials mapping
            oauth_credentials = data.get("oauth_credentials", [])
            for oauth_credential in oauth_credentials:
                provider = oauth_credential.get("provider")
                credentials = oauth_credential.get("credentials") or {}
                access_token = credentials.get("access_token")
                refresh_token = credentials.get("refresh_token")
                expires_at_in_seconds = credentials.get("expires_at_in_seconds")

                is_access_key_id_different = False

                if provider and access_token:
                    existing_provider = self.oauth_credentials_mapping.get(provider)
                    if existing_provider:
                        is_access_key_id_different = existing_provider.access_token != access_token
                        existing_provider.access_token = access_token
                        existing_provider.refresh_token = credentials.get("refresh_token")
                        existing_provider.expires_at_in_seconds = credentials.get("expires_at_in_seconds")
                    else:
                        is_access_key_id_different = True
                        self.oauth_credentials_mapping[provider] = BlocksOAuthCredentialsValues(**credentials)

                    if is_access_key_id_different:
                        print(f"Updating oauth credentials mapping for provider: {provider} with access_token: {access_token}, refresh_token: {refresh_token}, expires_at_in_seconds: {expires_at_in_seconds}")

                        # Ensure Github Token is kept up to date
                        if provider == BLOCKS_EVENT__PROVIDER.GITHUB.value:
                            os.environ["GITHUB_TOKEN"] = access_token
                            os.environ["GH_TOKEN"] = access_token
                            patch_blocks_runtime_config({
                                "GITHUB_TOKEN": access_token
                            })
                            update_git_credential("github.com", "oauth2", access_token)

                        if provider == BLOCKS_EVENT__PROVIDER.GITLAB.value:
                            os.environ["GITLAB_TOKEN"] = access_token
                            os.environ["GLAB_TOKEN"] = access_token
                            patch_blocks_runtime_config({
                                "GITLAB_TOKEN": access_token
                            })
                            update_git_credential("gitlab.com", "oauth2", access_token)

                    if provider == BLOCKS_EVENT__PROVIDER.CLAUDE.value and is_access_key_id_different:
                        claude_credentials_file = Path.home() / ".claude" / ".credentials.json"
                        claude_credentials_file.parent.mkdir(parents=True, exist_ok=True)
                        with open(claude_credentials_file.absolute(), "w") as f:
                            json.dump({
                                "claudeAiOauth": {
                                    "accessToken": access_token,
                                    "refreshToken": refresh_token,
                                    "expiresAt": expires_at_in_seconds * 1000, # convert to milliseconds as claude expects it
                                    "scopes": [
                                    "user:inference",
                                    "user:profile"
                                    ],
                                    "isMax": True
                                }
                            }, f)

                        print("===> CLAUDE CREDENTIALS CREATED/UPDATED <===")
                        with open(claude_credentials_file.absolute(), "r") as f:
                            print(f.read())
                        print("===> END OF CLAUDE CREDENTIALS CREATED/UPDATED <===")

                    if provider == BLOCKS_EVENT__PROVIDER.CODEX.value and is_access_key_id_different:
                        from datetime import datetime
                        codex_auth_file = Path.home() / ".codex" / "auth.json"
                        codex_auth_file.parent.mkdir(parents=True, exist_ok=True)

                        # Get additional fields from credentials
                        id_token = credentials.get("id_token")
                        account_id = credentials.get("account_id")

                        with open(codex_auth_file.absolute(), "w") as f:
                            json.dump({
                                "OPENAI_API_KEY": None,
                                "tokens": {
                                    "id_token": id_token,
                                    "access_token": access_token,
                                    "refresh_token": refresh_token,
                                    "account_id": account_id
                                },
                                "last_refresh": datetime.utcnow().isoformat() + "Z"
                            }, f, indent=2)

                        print("===> CODEX CREDENTIALS CREATED/UPDATED <===")
                        with open(codex_auth_file.absolute(), "r") as f:
                            print(f.read())
                        print("===> END OF CODEX CREDENTIALS CREATED/UPDATED <===")

            # Update the messages
            items = data.get("items", [])
            if len(items) > 0:
                # Track the highest timestamp & convert to BlocksMessage
                for item in items:
                    message = BlocksMessage(**item)
                    if message.ts > highest_ts:
                        highest_ts = message.ts
                    all_messages.append(message)
                page += 1
            else:
                break
        
        # Update the last processed timestamp if we found any messages
        if highest_ts > self.last_processed_ts:
            self.last_processed_ts = highest_ts
                
        return all_messages

    def format_message_history(self, messages: List[BlocksMessage]) -> str:
        """
        Format in xml as <user>...</user> or <assistant>...</assistant>
        """
        formatted_messages = []
        for message in messages:
            if message.role == "user":
                formatted_messages.append(f"<user>{message.message}</user>")
            elif message.role == "assistant":
                formatted_messages.append(f"<assistant>{message.message}</assistant>")
        return "\n".join(formatted_messages)
