from datetime import datetime, timedelta
import random
from marzban import MarzbanAPI, NodeResponse, ProxyHost, SystemStats, UserCreate, ProxySettings, UserResponse, ProxyInbound, UserModify
from database.models import VPNSubscription
import os
from typing import Optional, Dict, List
from zoneinfo import ZoneInfo
from dotenv import load_dotenv

load_dotenv()

class MarzbanService:
    def __init__(self):
        self.api = self._initialize_api()
        self._token = None
        self.countries = {}
        self.protocols = ["vless"]
        
    def _initialize_api(self) -> MarzbanAPI:
        """Initialize Marzban API with configuration from environment"""
        if os.getenv("ssh_on") == "true":
            return MarzbanAPI(
                base_url="http://127.0.0.1",
                ssh_username=os.getenv("ssh_login"),
                ssh_host=os.getenv("ssh_host"),
                ssh_port=os.getenv("ssh_port"),
                ssh_password=os.getenv("ssh_pass"),
                local_bind_host=os.getenv("local_bind_host"),
                local_bind_port=os.getenv("local_bind_port"),
                remote_bind_host=os.getenv("remote_bind_host"),
                remote_bind_port=os.getenv("remote_bind_port")
            )
        return MarzbanAPI(
            base_url=os.getenv("marzban_url"),
        )

    async def _get_token(self) -> str:
        """Get or refresh Marzban API token"""
        if not self._token or not hasattr(self, '_token_expires') or datetime.now() >= self._token_expires:
            token_response = await self.api.get_token(
                username=os.getenv("marzban_login"),
                password=os.getenv("marzban_pass")
            )
            self._token = token_response.access_token
            self._token_expires = datetime.now() + timedelta(hours=6)
        return self._token

    def _create_proxy_config(self, username: str, duration_days: int, country: str) -> Optional[UserCreate]:
        """Create proxy configuration for new user"""
        expire = datetime.now(ZoneInfo("Europe/Moscow")) + timedelta(days=duration_days)
        expire = int(expire.timestamp())
        proxies = {}
        inbounds = {}
        for protocol in self.protocols:
            # Check exists protocol in country
            if protocol in self.countries[country]:
                if protocol == "vless":
                    proxies[protocol] = ProxySettings(flow="xtls-rprx-vision")
                    # Get random inbound from country
                    inbound = random.choice(self.countries[country][protocol])
                    inbounds[protocol] = [inbound]
                else:
                    proxies[protocol] = ProxySettings()
                    # Get random inbound from country
                    inbound = random.choice(self.countries[country][protocol])
                    inbounds[protocol] = [inbound]
        if len(inbounds) == 0:
            return None
        return UserCreate(
            username=username,
            proxies=proxies,
            expire=expire,
            inbounds=inbounds
        )

    async def create_subscription(self, user_id: int, duration_days: int, sub_id: int, country: str) -> UserResponse:
        """Create new VPN subscription"""
        try:
            token = await self._get_token()
            username = f"{user_id}_{sub_id}"
            
            user_config = self._create_proxy_config(username, duration_days, country)
            if user_config is None:
                return None
            added_user = await self.api.add_user(user=user_config, token=token)
            
            print(f"Added user: {added_user}")
            return added_user
        except Exception as e:
            print(f"Error creating subscription: {e}")
            return None

    async def extend_subscription(self, user_id: int, sub_id: int, expire: int) -> UserResponse:
        """Extend VPN subscription"""
        try:
            token = await self._get_token()
            username = f"{user_id}_{sub_id}"
            user = UserModify(expire=expire)

            return await self.api.modify_user(username=username, user=user, token=token)
        except Exception as e:
            print(f"Error extending subscription: {e}")
            return None

    async def get_system_stats(self) -> SystemStats:
        """Get system statistics"""
        try:
            token = await self._get_token()
            return await self.api.get_system_stats(token=token)
        except Exception as e:
            print(f"Error getting system stats: {e}")
            return None

    async def get_system_state(self) -> Dict:
        """Get comprehensive system state including stats and node information"""
        try:
            stats = await self.get_system_stats()
            return stats   
        except Exception as e:
            print(f"Error getting system state: {e}")
            return None
        
    async def get_nodes(self) -> List[NodeResponse]:
        """Get list of all nodes"""
        try:
            token = await self._get_token()
            return await self.api.get_nodes(token=token)
        except Exception as e:
            print(f"Error getting nodes: {e}")
            return []

    async def get_node_settings(self) -> Dict:
        """Get node settings"""
        try:
            token = await self._get_token()
            return await self.api.get_node_settings(token=token)
        except Exception as e:
            print(f"Error getting node settings: {e}")
            return {}

    async def get_hosts(self) -> Dict[str, List[ProxyHost]]:
        """Get all proxy hosts grouped by inbound tag"""
        try:
            token = await self._get_token()
            return await self.api.get_hosts(token=token)
        except Exception as e:
            print(f"Error getting hosts: {e}")
            return {}
    
    async def get_inbounds(self) -> Dict:
        """Get all inbounds"""
        try:
            token = await self._get_token()
            return await self.api.get_inbounds(token=token)
        except Exception as e:
            print(f"Error getting inbounds: {e}")
            return {}
        
    async def get_countries(self) -> Dict[str, Dict[str, List[ProxyInbound]]]:
        """Get list of available countries from hosts"""
        inbounds = await self.get_inbounds()
        print(inbounds)
        for protocol in inbounds:
            if protocol in self.protocols:
                for inbound in inbounds[protocol]:
                    country = inbound["tag"].split("_")[0]
                    if country not in self.countries:
                        self.countries[country] = {}
                    if protocol not in self.countries[country]:
                        self.countries[country][protocol] = []
                    self.countries[country][protocol].append(inbound["tag"])
        return self.countries
        # try:
        #     inbounds = await self.get_inbounds()
        #     print(inbounds)
        #     for protocol in inbounds:
        #         if protocol in self.protocols:
        #             for inbound in inbounds[protocol]:
        #                 country = inbound["tag"].split("_")[0]
        #                 if country not in self.countries:
        #                     if protocol not in self.countries[country]:
        #                         self.countries[country] = {protocol: []}
        #                     self.countries[country][protocol].append(inbound["tag"])
        #     return self.countries
        # except Exception as e:
        #     print(f"Error getting countries: {e}")
        #     return {}

    async def close(self):
        """Close API connection"""
        if self.api:
            await self.api.close()


# Create global instance
marzban_service = MarzbanService()