#!/usr/bin/env python3
"""
Claude Code Status Bar Monitor - Final Fixed Version
Resolves dependency issues, ensuring operation in any environment
"""

import json
import sys
import logging
import os
import subprocess
import shutil
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Dict, List, Optional, Any

# Suppress log output
logging.basicConfig(level=logging.ERROR)

# ANSI color codes
class Colors:
    GREEN = '\033[32m'
    YELLOW = '\033[33m'
    RED = '\033[31m'
    RESET = '\033[0m'

def try_original_analysis() -> Optional[Dict[str, Any]]:
    """Try to use the installed claude-monitor package"""
    try:
        
        # Check if claude-monitor is installed
        claude_monitor_cmd = shutil.which('claude-monitor')
        if not claude_monitor_cmd:
            # Try other command aliases
            for cmd in ['cmonitor', 'ccmonitor', 'ccm']:
                claude_monitor_cmd = shutil.which(cmd)
                if claude_monitor_cmd:
                    break
        
        if not claude_monitor_cmd:
            logging.info("claude-monitor not found. Install with: uv tool install claude-monitor")
            return None
        
        # Find the Python interpreter used by claude-monitor
        # Check common installation paths
        possible_paths = [
            Path.home() / ".local/share/uv/tools/claude-monitor/bin/python",
            Path.home() / ".uv/tools/claude-monitor/bin/python",
            Path.home() / ".local/pipx/venvs/claude-monitor/bin/python",  # pipx installation
        ]
        
        claude_python = None
        for path in possible_paths:
            if path.exists():
                claude_python = str(path)
                break
        
        if not claude_python:
            # Try to extract from the shebang of claude-monitor script
            try:
                with open(claude_monitor_cmd, 'r') as f:
                    first_line = f.readline()
                    if first_line.startswith('#!'):
                        claude_python = first_line[2:].strip()
            except:
                pass
        
        if not claude_python:
            logging.info("Could not find claude-monitor Python interpreter")
            return None
        
        # Use subprocess to run analysis with the correct Python
        code = """
import json
import sys
try:
    # Version compatibility check
    import claude_monitor
    version = getattr(claude_monitor, '__version__', 'unknown')
    
    from claude_monitor.data.analysis import analyze_usage
    from claude_monitor.core.plans import get_token_limit
    
    result = analyze_usage(hours_back=192, quick_start=False)
    blocks = result.get('blocks', [])
    
    if not blocks:
        print(json.dumps(None))
        sys.exit(0)
    
    # Get active sessions
    active_blocks = [b for b in blocks if b.get('isActive', False)]
    if not active_blocks:
        print(json.dumps(None))
        sys.exit(0)
    
    current_block = active_blocks[0]
    
    # Get P90 limit with compatibility handling
    try:
        token_limit = get_token_limit('custom', blocks)
    except TypeError:
        # Try old API signature
        try:
            token_limit = get_token_limit('custom')
        except:
            token_limit = 113505
    except:
        token_limit = 113505
    
    # Calculate dynamic cost limit using P90 method similar to claude-monitor
    try:
        # Get all historical costs from blocks for P90 calculation
        all_costs = []
        for block in blocks:
            cost = block.get('costUSD', 0)
            if cost > 0:
                all_costs.append(cost)
        
        # Also collect message counts for P90 calculation
        all_messages = []
        for block in blocks:
            msg_count = block.get('sentMessagesCount', len(block.get('entries', [])))
            if msg_count > 0:
                all_messages.append(msg_count)
        
        if len(all_costs) >= 5:
            # Use P90 calculation similar to claude-monitor
            all_costs.sort()
            all_messages.sort()
            p90_index = int(len(all_costs) * 0.9)
            p90_cost = all_costs[min(p90_index, len(all_costs) - 1)]
            # Calculate message limit using P90 method
            if all_messages:
                p90_msg_index = int(len(all_messages) * 0.9)
                p90_messages = all_messages[min(p90_msg_index, len(all_messages) - 1)]
                message_limit = max(int(p90_messages * 1.2), 100)  # Similar to cost calculation
            else:
                message_limit = 755  # Default based on your example
            
            # Apply similar logic to claude-monitor (seems to use a different multiplier)
            cost_limit = max(p90_cost * 1.004, 50.0)  # Adjusted to match observed behavior
        else:
            # Fallback to static limit
            from claude_monitor.core.plans import get_cost_limit
            cost_limit = get_cost_limit('custom')
            message_limit = 755  # Default
    except:
        cost_limit = 90.26  # fallback
    
    # Handle different field name conventions for compatibility
    total_tokens = (current_block.get('totalTokens', 0) or 
                   current_block.get('total_tokens', 0) or 0)
    cost_usd = (current_block.get('costUSD', 0.0) or 
               current_block.get('cost_usd', 0.0) or 
               current_block.get('cost', 0.0) or 0.0)
    entries = current_block.get('entries', []) or []
    messages_count = current_block.get('sentMessagesCount', len(entries))
    is_active = current_block.get('isActive', current_block.get('is_active', False))
    
    output = {
        'total_tokens': total_tokens,
        'token_limit': token_limit,
        'cost_usd': cost_usd,
        'cost_limit': cost_limit,
        'messages_count': messages_count,
        'message_limit': message_limit,
        'entries_count': len(entries),
        'is_active': is_active,
        'plan_type': 'CUSTOM',
        'source': 'original'
    }
    print(json.dumps(output))
except Exception as e:
    print(json.dumps(None))
    sys.exit(1)
"""
        
        # Run the code with the claude-monitor Python interpreter
        result = subprocess.run(
            [claude_python, '-c', code],
            capture_output=True,
            text=True,
            timeout=10
        )
        
        if result.returncode == 0 and result.stdout:
            data = json.loads(result.stdout.strip())
            if data:
                return data
        
        return None
        
    except Exception as e:
        logging.error(f"Original analysis failed: {e}")
        return None

def direct_data_analysis() -> Optional[Dict[str, Any]]:
    """Directly analyze Claude data files, completely independent implementation"""
    try:
        def build_candidate_paths() -> List[Path]:
            """Collect plausible data directories in priority order."""
            paths: List[Path] = []
            
            # Respect Claude Code env override
            env_dir = os.environ.get("CLAUDE_CONFIG_DIR")
            if env_dir:
                env_path = Path(env_dir).expanduser()
                if env_path.name == ".claude":
                    paths.append(env_path)
                    paths.append(env_path / "projects")
                else:
                    paths.append(env_path / ".claude")
                    paths.append(env_path / ".claude" / "projects")
            
            # Running from inside .claude
            cwd = Path.cwd()
            if cwd.name == ".claude":
                paths.append(cwd)
                paths.append(cwd / "projects")
            
            # Standard locations
            paths.extend([
                Path.home() / '.claude' / 'projects',
                Path.home() / '.config' / 'claude' / 'projects',
                Path.home() / '.claude',
            ])
            
            # Deduplicate while preserving order
            seen = set()
            unique_paths: List[Path] = []
            for p in paths:
                if p not in seen:
                    unique_paths.append(p)
                    seen.add(p)
            return unique_paths

        data_path = None
        for path in build_candidate_paths():
            if path.exists() and path.is_dir():
                data_path = path
                break
        
        if not data_path:
            return None
        
        # Collect data from the last 5 hours (simulate session window)
        cutoff_time = datetime.now(timezone.utc) - timedelta(hours=5)
        current_session_data = []
        
        # Collect historical data for P90 calculation
        history_cutoff = datetime.now(timezone.utc) - timedelta(days=8)
        all_sessions = []
        current_session_tokens = 0
        current_session_cost = 0.0
        last_time = None
        
        # Read all JSONL files
        for jsonl_file in sorted(data_path.rglob("*.jsonl"), key=lambda f: f.stat().st_mtime):
            try:
                with open(jsonl_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue
                        
                        try:
                            data = json.loads(line)
                            
                            # Parse timestamp
                            timestamp_str = data.get('timestamp', '')
                            if not timestamp_str:
                                continue
                            
                            if timestamp_str.endswith('Z'):
                                timestamp_str = timestamp_str[:-1] + '+00:00'
                            
                            timestamp = datetime.fromisoformat(timestamp_str)
                            
                            # Extract usage data
                            usage = data.get('usage', {})
                            if not usage and 'message' in data and isinstance(data['message'], dict):
                                usage = data['message'].get('usage', {})
                            
                            if not usage:
                                continue
                            
                            # Calculate tokens
                            input_tokens = usage.get('input_tokens', 0)
                            output_tokens = usage.get('output_tokens', 0)
                            cache_creation = usage.get('cache_creation_input_tokens', 0)
                            cache_read = usage.get('cache_read_input_tokens', 0)
                            
                            total_tokens = input_tokens + output_tokens + cache_creation
                            
                            if total_tokens == 0:
                                continue
                            
                            # Estimate cost (simplified pricing model)
                            # Based on Sonnet 3.5 pricing: input $3/M tokens, output $15/M tokens
                            cost = (input_tokens * 3 + output_tokens * 15 + cache_creation * 3.75) / 1000000
                            
                            entry = {
                                'timestamp': timestamp,
                                'total_tokens': total_tokens,
                                'cost': cost,
                                'input_tokens': input_tokens,
                                'output_tokens': output_tokens,
                                'cache_creation': cache_creation,
                                'cache_read': cache_read
                            }
                            
                            # Current 5-hour session data
                            if timestamp >= cutoff_time:
                                current_session_data.append(entry)
                            
                            # Historical session grouping (for P90 calculation)
                            if timestamp >= history_cutoff:
                                if (last_time is None or 
                                    (timestamp - last_time).total_seconds() > 5 * 3600):
                                    # Save previous session
                                    if current_session_tokens > 0:
                                        all_sessions.append({
                                            'tokens': current_session_tokens,
                                            'cost': current_session_cost
                                        })
                                    # Start new session
                                    current_session_tokens = total_tokens
                                    current_session_cost = cost
                                else:
                                    # Continue current session
                                    current_session_tokens += total_tokens
                                    current_session_cost += cost
                                
                                last_time = timestamp
                        
                        except (json.JSONDecodeError, ValueError, TypeError):
                            continue
                            
            except Exception:
                continue
        
        # Save last session
        if current_session_tokens > 0:
            all_sessions.append({
                'tokens': current_session_tokens,
                'cost': current_session_cost
            })
        
        if not current_session_data:
            return None
        
        # Calculate current session statistics
        total_tokens = sum(e['total_tokens'] for e in current_session_data)
        total_cost = sum(e['cost'] for e in current_session_data)
        
        # Calculate P90 limit
        if len(all_sessions) >= 5:
            session_tokens = [s['tokens'] for s in all_sessions]
            session_costs = [s['cost'] for s in all_sessions]
            session_tokens.sort()
            session_costs.sort()
            
            p90_index = int(len(session_tokens) * 0.9)
            token_limit = max(session_tokens[min(p90_index, len(session_tokens) - 1)], 19000)
            cost_limit = max(session_costs[min(p90_index, len(session_costs) - 1)] * 1.2, 18.0)
        else:
            # Default limits
            if total_tokens > 100000:
                token_limit, cost_limit = 220000, 140.0
            elif total_tokens > 50000:
                token_limit, cost_limit = 88000, 35.0
            else:
                token_limit, cost_limit = 19000, 18.0
        
        return {
            'total_tokens': total_tokens,
            'token_limit': int(token_limit),
            'cost_usd': total_cost,
            'cost_limit': cost_limit,
            'messages_count': len(current_session_data),  # Each entry is a message
            'message_limit': 755,  # Default fallback
            'entries_count': len(current_session_data),
            'is_active': True,
            'plan_type': 'CUSTOM' if len(all_sessions) >= 5 else 'AUTO',
            'source': 'direct'
        }
        
    except Exception as e:
        logging.error(f"Direct analysis failed: {e}")
        return None

def get_current_model() -> tuple[str, str]:
    """Get current Claude model from settings.json and display name from stdin"""
    try:
        # First, get model from settings.json
        settings_path = Path.home() / '.claude' / 'settings.json'
        model = "unknown"
        if settings_path.exists():
            with open(settings_path, 'r', encoding='utf-8') as f:
                settings = json.load(f)
                model = settings.get('model', 'unknown')
        
        # Then try to get display name from stdin (Claude Code provides this)
        display_name = "Unknown"
        if not sys.stdin.isatty():
            stdin_data = sys.stdin.read()
            if stdin_data:
                try:
                    claude_data = json.loads(stdin_data)
                    display_name = claude_data.get('model', {}).get('display_name', 'Unknown')
                except (json.JSONDecodeError, TypeError):
                    # If stdin parse fails, use model as display name
                    display_name = model.title() if model != 'unknown' else 'Unknown'
        else:
            # No stdin, use model as display name
            display_name = model.title() if model != 'unknown' else 'Unknown'
        
        return model, display_name
            
    except Exception:
        return "unknown", "Unknown"

def calculate_reset_time(reset_hour: Optional[int] = None) -> str:
    """Calculate time until session reset (5-hour rolling window or custom hour)"""
    # If user pins a reset hour, honor it before any external calls
    if reset_hour is not None and 0 <= reset_hour <= 23:
        now = datetime.now()
        target = now.replace(hour=reset_hour, minute=0, second=0, microsecond=0)
        if target <= now:
            target = target + timedelta(days=1)
        diff = target - now
        total_minutes = int(diff.total_seconds() / 60)
        hours = total_minutes // 60
        mins = total_minutes % 60
        return f"{hours}h {mins:02d}m"

    try:
        
        # Try the same method as try_original_analysis to get session data
        claude_monitor_cmd = shutil.which('claude-monitor')
        if claude_monitor_cmd:
            # Find Python interpreter
            possible_paths = [
                Path.home() / ".local/share/uv/tools/claude-monitor/bin/python",
                Path.home() / ".uv/tools/claude-monitor/bin/python",
                Path.home() / ".local/pipx/venvs/claude-monitor/bin/python",  # pipx installation
            ]
            
            claude_python = None
            for path in possible_paths:
                if path.exists():
                    claude_python = str(path)
                    break
            
            if not claude_python:
                try:
                    with open(claude_monitor_cmd, 'r') as f:
                        first_line = f.readline()
                        if first_line.startswith('#!'):
                            claude_python = first_line[2:].strip()
                except:
                    pass
            
            if claude_python:
                code = """
import json
from datetime import datetime, timedelta, timezone
try:
    from claude_monitor.data.analysis import analyze_usage
    
    result = analyze_usage(hours_back=192, quick_start=False)
    blocks = result.get('blocks', [])
    
    if blocks:
        active_blocks = [b for b in blocks if b.get('isActive', False)]
        if active_blocks:
            current_block = active_blocks[0]
            start_time = current_block.get('startTime')
            
            if start_time:
                # Parse start time
                if isinstance(start_time, str):
                    if start_time.endswith('Z'):
                        start_time = start_time[:-1] + '+00:00'
                    session_start = datetime.fromisoformat(start_time)
                else:
                    session_start = start_time
                
                # Session lasts 5 hours
                session_end = session_start + timedelta(hours=5)
                now = datetime.now(timezone.utc)
                
                if session_end > now:
                    diff = session_end - now
                    total_minutes = int(diff.total_seconds() / 60)
                    
                    if total_minutes > 60:
                        hours = total_minutes // 60
                        mins = total_minutes % 60
                        print(f"{hours}h {mins:02d}m")
                    else:
                        print(f"{total_minutes}m")
                    import sys
                    sys.exit(0)
except:
    pass
print("")
"""
                result = subprocess.run(
                    [claude_python, '-c', code],
                    capture_output=True,
                    text=True,
                    timeout=5
                )
                
                if result.returncode == 0 and result.stdout.strip():
                    return result.stdout.strip()
    except:
        pass
    
    # Fallback: estimate reset time (assume session started within the last 5 hours)
    now = datetime.now()
    # Assume reset time is 2 PM (consistent with original project display)
    today_2pm = now.replace(hour=14, minute=0, second=0, microsecond=0)
    tomorrow_2pm = today_2pm + timedelta(days=1)
    
    # Choose next 2 PM
    next_reset = tomorrow_2pm if now >= today_2pm else today_2pm
    diff = next_reset - now
    
    total_minutes = int(diff.total_seconds() / 60)
    hours = total_minutes // 60
    mins = total_minutes % 60
    
    return f"{hours}h {mins:02d}m"

def format_number(num: float) -> str:
    """Format number display"""
    if num >= 1000000:
        return f"{num/1000000:.1f}M"
    elif num >= 1000:
        return f"{num/1000:.1f}k"
    else:
        return f"{num:.0f}"

PLAN_PRESETS = {
    # Basic Anthropic-style fallbacks
    "max5": {"token_limit": 88000, "cost_limit": 35.0, "message_limit": 500},
    "max20": {"token_limit": 220000, "cost_limit": 140.0, "message_limit": 900},
    "pro": {"token_limit": 19000, "cost_limit": 18.0, "message_limit": 755},
    # z.ai subscription estimates (prompt counts mapped to message_limit)
    "zai-lite": {"token_limit": 400000, "cost_limit": 50.0, "message_limit": 120},
    "zai-pro": {"token_limit": 2000000, "cost_limit": 150.0, "message_limit": 600},
    "zai-max": {"token_limit": 8000000, "cost_limit": 600.0, "message_limit": 2400},
}

def apply_plan_override(usage_data: Dict[str, Any], plan_name: Optional[str]) -> Dict[str, Any]:
    """Apply plan override to usage data when user specifies a plan."""
    if not plan_name:
        return usage_data
    
    normalized = plan_name.lower().replace("_", "-").replace(".", "-")
    preset = PLAN_PRESETS.get(normalized)
    
    if preset:
        usage_data = usage_data.copy()
        usage_data['token_limit'] = preset['token_limit']
        usage_data['cost_limit'] = preset['cost_limit']
        usage_data['message_limit'] = preset['message_limit']
        usage_data['plan_type'] = normalized
    else:
        usage_data = usage_data.copy()
        usage_data['plan_type'] = normalized
    
    return usage_data

def generate_statusbar_text(usage_data: Dict[str, Any], reset_hour: Optional[int] = None) -> str:
    """Generate status bar text"""
    total_tokens = usage_data['total_tokens']
    token_limit = usage_data['token_limit']
    cost_usd = usage_data['cost_usd']
    cost_limit = usage_data['cost_limit']
    messages_count = usage_data.get('messages_count', 0)
    message_limit = usage_data.get('message_limit', 755)
    plan_type = usage_data['plan_type']
    source = usage_data.get('source', 'unknown')
    
    # Calculate percentage - use cost as primary metric
    cost_percentage = (cost_usd / cost_limit) * 100 if cost_limit > 0 else 0
    token_percentage = (total_tokens / token_limit) * 100 if token_limit > 0 else 0
    message_percentage = (messages_count / message_limit) * 100 if message_limit > 0 else 0
    
    # Choose color based on cost percentage (main metric)
    if cost_percentage >= 90:
        color = Colors.RED
    elif cost_percentage >= 70:
        color = Colors.RED
    elif cost_percentage >= 30:
        color = Colors.YELLOW
    else:
        color = Colors.GREEN
    
    # Messages display without percentage
    messages_display = f'📨:{messages_count}/{message_limit}'
    
    # Get current model and reset time
    current_model, display_name = get_current_model()
    reset_time = calculate_reset_time(reset_hour=reset_hour)
    
    # Format text - integrated model display
    tokens_text = f"🔋:{format_number(total_tokens)}/{format_number(token_limit)}"
    cost_text = f"💰:{cost_usd:.2f}/{cost_limit:.2f}"
    time_text = f"⌛️:{reset_time}"
    model_text = f"🤖:{current_model}({display_name})"
    
    status_text = f"{tokens_text} | {cost_text} | {messages_display} | {time_text} | {model_text}"
    
    return f"{color}{status_text}{Colors.RESET}"

def check_for_updates():
    """Check for updates once per day"""
    try:
        from datetime import datetime
        
        # Check if we should run update check
        last_check_file = Path.home() / '.claude-statusbar-last-check'
        now = datetime.now()
        
        should_check = True
        if last_check_file.exists():
            try:
                with open(last_check_file, 'r') as f:
                    last_check_str = f.read().strip()
                    last_check = datetime.fromisoformat(last_check_str)
                    # Check once per day
                    if (now - last_check).days < 1:
                        should_check = False
            except:
                pass
        
        if should_check:
            # Run update check in background
            from .updater import check_and_upgrade
            success, message = check_and_upgrade()
            
            # Update last check time
            with open(last_check_file, 'w') as f:
                f.write(now.isoformat())
            
            # If upgrade was successful, notify user
            if success:
                print(f"🔄 {message}", file=sys.stderr)
                
    except Exception:
        # Silently fail - don't interrupt main functionality
        pass

def build_json_output(usage_data: Dict[str, Any], reset_time: str, model: str, display_name: str) -> Dict[str, Any]:
    """Create machine-readable payload."""
    return {
        "success": True,
        "usage": {
            "total_tokens": usage_data.get("total_tokens", 0),
            "token_limit": usage_data.get("token_limit", 0),
            "cost_usd": usage_data.get("cost_usd", 0.0),
            "cost_limit": usage_data.get("cost_limit", 0.0),
            "messages_count": usage_data.get("messages_count", 0),
            "message_limit": usage_data.get("message_limit", 0),
            "plan_type": usage_data.get("plan_type"),
            "source": usage_data.get("source", "unknown"),
        },
        "meta": {
            "model": model,
            "display_name": display_name,
            "reset_time": reset_time,
            "generated_at": datetime.now(timezone.utc).isoformat(),
        },
    }

def main(json_output: bool = False, plan: Optional[str] = None, reset_hour: Optional[int] = None):
    """Main function"""
    try:
        # Check for updates (silent, once per day)
        check_for_updates()
        
        # First try original project analysis
        usage_data = try_original_analysis()
        
        # If failed, use direct analysis
        if not usage_data:
            usage_data = direct_data_analysis()

        # Apply plan override if requested
        if usage_data:
            usage_data = apply_plan_override(usage_data, plan)
        
        if not usage_data:
            # Final fallback
            reset_time = calculate_reset_time(reset_hour=reset_hour)
            current_model, display_name = get_current_model()
            fallback_usage = {
                "total_tokens": 0,
                "token_limit": PLAN_PRESETS.get("pro", {}).get("token_limit", 19000),
                "cost_usd": 0.0,
                "cost_limit": PLAN_PRESETS.get("pro", {}).get("cost_limit", 18.0),
                "messages_count": 0,
                "message_limit": PLAN_PRESETS.get("pro", {}).get("message_limit", 755),
                "plan_type": plan or "unknown",
                "source": "fallback",
            }
            fallback_usage = apply_plan_override(fallback_usage, plan)
            if json_output:
                payload = build_json_output(fallback_usage, reset_time, current_model, display_name)
                payload["success"] = False
                payload["error"] = "No usage data found"
                print(json.dumps(payload))
            else:
                tokens_text = f"🔋:0/{format_number(fallback_usage['token_limit'])}"
                cost_text = f"💰:0.00/{fallback_usage['cost_limit']:.2f}"
                messages_text = f"📨:0/{fallback_usage['message_limit']}"
                print(f"{tokens_text} | {cost_text} | {messages_text} | ⌛️:{reset_time} | 🤖:{current_model}({display_name})")
            return
        
        # Generate status bar text
        if json_output:
            current_model, display_name = get_current_model()
            reset_time = calculate_reset_time(reset_hour=reset_hour)
            payload = build_json_output(usage_data, reset_time, current_model, display_name)
            print(json.dumps(payload))
        else:
            status_text = generate_statusbar_text(usage_data, reset_hour=reset_hour)
            print(status_text)
        
    except Exception as e:
        # Basic display on error
        reset_time = calculate_reset_time(reset_hour=reset_hour)
        current_model, display_name = get_current_model()
        if json_output:
            payload = {
                "success": False,
                "error": str(e),
                "meta": {
                    "model": current_model,
                    "display_name": display_name,
                    "reset_time": reset_time,
                    "generated_at": datetime.now(timezone.utc).isoformat(),
                },
            }
            print(json.dumps(payload))
        else:
            print(f"🔋:0/19k | 💰:0.00/18.00 | 📨:0/755 | ⌛️:{reset_time} | 🤖:{current_model}({display_name}) | ❌")

if __name__ == '__main__':
    main()
