import asyncio
import os
import json
import traceback
from typing import Any, Dict, List, Optional
import requests
from requests.auth import HTTPBasicAuth

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource
from dotenv import load_dotenv

# Add current directory to path to ensure imports work if run from everywhere
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.billing import BillingDatabase

# Load environment variables
load_dotenv()

# Initialize Server
app = Server("billing-mcp")

# Global instance
billing_db: Optional[BillingDatabase] = None
config: Dict[str, Any] = {}

# Map to store which router belongs to which ISP (Billing DB Server ID)
# Format: { "router_host_or_name": "isp_server_id" }
router_isp_map: Dict[str, str] = {}
# Map to store resolved router config
# Format: { "router_host_or_name": { "host": "...", "port": ..., "user": "...", "pass": "..." } }
resolved_router_configs: Dict[str, Dict] = {}

def load_config():
    global config, router_isp_map, resolved_router_configs
    config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'config.json')
    if os.path.exists(config_path):
        with open(config_path, 'r') as f:
            config = json.load(f)
            
        # Parse ISP configs to build router map
        if 'isps' in config:
            for isp_id, isp_data in config['isps'].items():
                routers = isp_data.get('routers', {})
                for r_name, r_conf in routers.items():
                    # Map Name -> ISP ID
                    router_isp_map[r_name] = isp_id
                    # Map IP -> ISP ID
                    if 'host' in r_conf:
                        router_isp_map[r_conf['host']] = isp_id
                    
                    # Store resolved config for easy lookup
                    resolved_router_configs[r_name] = r_conf
                    if 'host' in r_conf:
                        resolved_router_configs[r_conf['host']] = r_conf
    else:
        print(f"Warning: Config file not found at {config_path}", file=sys.stderr)

def get_billing_db() -> BillingDatabase:
    global billing_db
    if billing_db is None:
        load_config()
        
        # Prepare configs for BillingDatabase
        db_configs = {}
        
        if 'isps' in config:
            # Flatten ISP structure for BillingDatabase
            # BillingDatabase expects { "server_id": { config... } }
            for isp_id, isp_data in config['isps'].items():
                if 'billing' in isp_data:
                    # Provide default alias if missing
                    billing_conf = isp_data['billing']
                    if 'alias' not in billing_conf and 'alias' in isp_data:
                         billing_conf['alias'] = isp_data['alias']
                    db_configs[isp_id] = billing_conf
        elif 'billing_databases' in config:
            # Legacy fallback
            db_configs = config['billing_databases']
            
        billing_db = BillingDatabase(db_configs) 
        
    return billing_db

@app.list_tools()
async def list_tools() -> List[Tool]:
    return [
        Tool(
            name="search_customers",
            description="Search for customers in the billing database. Returns specific customer details only.",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Search query (User ID, Name, or IP address). Leave empty to list all (use with limit)."
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum number of results to return. Default is 5.",
                        "default": 5
                    },
                    "isp_name": {
                        "type": "string",
                        "description": "Optional: Specific ISP/Provider name to search (e.g. 'dimensi'). If omitted, defaults to first available."
                    }
                },
                "required": ["query"]
            }
        ),
        Tool(
            name="refresh_billing_snapshot",
            description="Force refresh of the local billing snapshot from the MySQL database.",
            inputSchema={
                "type": "object",
                "properties": {
                     "isp_name": {
                        "type": "string",
                        "description": "Optional: Specific ISP/Provider name to refresh. If omitted, refreshes default."
                    }
                }
            }
        ),
        Tool(
             name="audit_router_users",
             description="Audit a MikroTik router to find users that are NOT in the billing database. Credentials will be auto-loaded from config if not provided.",
             inputSchema={
                "type": "object",
                "properties": {
                    "router_host": {
                        "type": "string",
                        "description": "IP address or Name of the router (e.g. 'ccr1036')"
                    },
                    "router_port": {
                        "type": "integer",
                        "description": "Port of the router API (default 80)",
                        "default": 80
                    },
                     "router_user": {
                        "type": "string",
                        "description": "Username for router (optional if configured)"
                    },
                     "router_pass": {
                        "type": "string",
                        "description": "Password for router (optional if configured)"
                    }
                },
                "required": ["router_host"]
            }
        ),
        Tool(
            name="list_system_info",
            description="List all configured ISPs and their associated Routers.",
            inputSchema={
                "type": "object",
                "properties": {},
            }
        )
    ]

@app.call_tool()
async def call_tool(name: str, arguments: Any) -> List[TextContent | ImageContent | EmbeddedResource]:
    db = get_billing_db()
    
    if name == "search_customers":
        query = arguments.get("query", "")
        limit = arguments.get("limit", 5)
        isp_name = arguments.get("isp_name", "1") # Default logic in simple mode
        
        # If isp_name is not provided, BillingDatabase likely handles "1" as default fallback 
        # or we might want to iterate all if query is global? 
        # For now, let's stick to single server/default behavior unless specified.
        
        result = db.search_customers(query, server_id=isp_name, limit=limit)
        
        if not result['success']:
            return [TextContent(type="text", text=f"Error searching ({isp_name}): {result.get('error')}")]
        
        customers = result.get('customers', [])
        if not customers:
            return [TextContent(type="text", text=f"No customers found matching '{query}'.")]
            
        # Format output
        output_lines = [f"Found {len(customers)} customers (ISP: {result.get('server_id')}):"]
        for c in customers:
            # Format key details
            details = [
                f"User: {c.get('user_mikrotik', 'N/A')}",
                f"Name: {c.get('name', 'N/A')}",
                f"Status: {c.get('c_status', 'N/A')}",
                f"Address: {c.get('address', 'N/A')}",
                f"Packet: {c.get('user_profile', 'N/A')}"
            ]
            output_lines.append(" | ".join(details))
            
        return [TextContent(type="text", text="\n".join(output_lines))]

    elif name == "refresh_billing_snapshot":
        isp_name = arguments.get("isp_name", "1")
        success = db.refresh_snapshot(server_id=isp_name)
        if success:
            return [TextContent(type="text", text=f"Billing snapshot for '{isp_name}' refreshed successfully.")]
        else:
            return [TextContent(type="text", text="Failed to refresh billing snapshot. Check logs.")]

    elif name == "audit_router_users":
        host_arg = arguments.get("router_host")
        port = arguments.get("router_port", 80)
        user = arguments.get("router_user")
        password = arguments.get("router_pass")
        
        # 0. Identify ISP/Provider context
        target_host = host_arg
        isp_context = "1" # Default
        
        # Lookup in map
        if host_arg in router_isp_map:
            isp_context = router_isp_map[host_arg]
        
        # Auto-lookup credentials
        if not user or not password:
            found_config = resolved_router_configs.get(host_arg)
            
            # Fallback check internal routers config (legacy)
            if not found_config and 'routers' in config:
                if host_arg in config['routers']:
                    found_config = config['routers'][host_arg]
                else:
                    for k, v in config['routers'].items():
                         if v.get('host') == host_arg:
                            found_config = v
                            break

            if found_config:
                user = found_config.get('user')
                password = found_config.get('pass')
                target_host = found_config.get('host', target_host)
                port = found_config.get('port', port)
            else:
                 return [TextContent(type="text", text=f"Error: Credentials not provided and router '{host_arg}' not found in config.")]

        # 1. Fetch router secrets
        url = f"http://{target_host}:{port}/rest/ppp/secret"
        try:
            resp = requests.get(url, auth=HTTPBasicAuth(user, password), timeout=10)
            resp.raise_for_status()
            secrets_data = resp.json()
            router_secrets = secrets_data if isinstance(secrets_data, list) else [secrets_data] if secrets_data else []
        except Exception as e:
             return [TextContent(type="text", text=f"Error connecting to router {target_host}: {str(e)}")]

        # 2. Get Billing Users (FROM SPECIFIC ISP CONTEXT)
        # We assume strict isolation: Only compare against the ISP's billing DB.
        res = db.search_customers("", server_id=isp_context, limit=10000)
        
        if not res['success']:
             return [TextContent(type="text", text=f"Error loading billing data for ISP '{isp_context}': {res.get('error')}")]
        
        billing_users = {c.get('user_mikrotik') for c in res['customers'] if c.get('user_mikrotik')}
        
        # 3. Compare
        unregistered = []
        for s in router_secrets:
            s_name = s.get('name')
            if s_name and s_name not in billing_users:
                unregistered.append(f"{s_name} (Profile: {s.get('profile', '?')})")
        
        if unregistered:
            return [TextContent(type="text", text=f"Found {len(unregistered)} Unregistered Users on {target_host} (ISP: {isp_context}):\n" + "\n".join(sorted(unregistered)))]
        else:
            return [TextContent(type="text", text=f"All {len(router_secrets)} users on router {target_host} (ISP: {isp_context}) are valid.")]

    elif name == "list_system_info":
        if 'isps' not in config:
            return [TextContent(type="text", text="No ISPs configured (Legacy mode or empty config).")]
            
        output = ["Computed System Configuration:", ""]
        
        for isp_id, data in config['isps'].items():
            alias = data.get('alias', isp_id)
            billing_host = data.get('billing', {}).get('host', 'N/A')
            billing_db_name = data.get('billing', {}).get('name', 'N/A')
            
            output.append(f"🌐 ISP: {alias} (ID: {isp_id})")
            output.append(f"   Using Billing DB: {billing_db_name} @ {billing_host}")
            
            routers = data.get('routers', {})
            if routers:
                output.append("   📡 Registered Routers:")
                for r_name, r_conf in routers.items():
                    host = r_conf.get('host', 'N/A')
                    port = r_conf.get('port', 80)
                    output.append(f"      - {r_name} ({host}:{port})")
            else:
                 output.append("   ⚠️ No routers configured.")
            
            output.append("") # Empty line separator
            
        return [TextContent(type="text", text="\n".join(output))]

    raise ValueError(f"Unknown tool: {name}")

async def main():
    # Run the server using stdin/stdout
    async with stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options()
        )

if __name__ == "__main__":
    asyncio.run(main())
