import json
import os
import stat
import requests
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, List, Tuple, Any

MODELS = {
    'gpt-3.5-turbo': {'context': 16385, 'endpoint': 'gpt-3.5-turbo'},
    'gpt-4-turbo': {'context': 128000, 'endpoint': 'gpt-4-turbo'},
    'gpt-4o': {'context': 128000, 'endpoint': 'gpt-4o'},
    'gpt-5': {'context': 256000, 'endpoint': 'gpt-5'},
    'gpt-5-nano': {'context': 32000, 'endpoint': 'gpt-5-nano'},
}

DEFAULT_MODEL = 'gpt-5'

OPENAI_API_URL = 'https://api.openai.com/v1/chat/completions'

SYSTEM_PROMPT_OPTIMIZE = '''You are a C++ expert specializing in pybind11 bindings and the IncludeCPP framework.

Rules:
1. Never remove existing functions or code blocks unless explicitly instructed
2. Preserve all public API signatures exactly
3. Focus on: performance, memory safety, pybind11 best practices
4. Use modern C++17 features where beneficial
5. Maintain namespace includecpp structure
6. Do not add comments, docstrings, or explanatory text to the code
7. Do not add AI-typical output markers or annotations

If a change would remove or significantly alter existing functionality, you MUST prefix that section with:
CONFIRM_REQUIRED: <brief description>

Output format for each file:
FILE: <relative_path>
CHANGES:
- Line X: <what changed> - <why>
- Line Y-Z: <what changed> - <why>
```cpp
<complete file content with no comments about changes>
```

Only output files that need changes. If no optimizations needed, respond with: NO_CHANGES_NEEDED'''

SYSTEM_PROMPT_FIX = '''You are a C++ expert analyzing code for the IncludeCPP framework.

Context:
- IncludeCPP generates pybind11 bindings from .cp plugin files
- SOURCE() and HEADER() directives link to C++ source files
- All exposed code must be in namespace includecpp
- Modules compile to .pyd (Windows) or .so (Linux/Mac) Python extensions

Analyze the provided code and fix issues. Rules:
1. Never remove functions or change public API signatures
2. Fix syntax errors, type issues, memory problems
3. Improve pybind11 compatibility
4. Do not add comments or docstrings
5. Preserve exact logic and behavior

If removal is necessary, prefix with:
CONFIRM_REQUIRED: <description>

Output format for each file:
FILE: <path>
CHANGES:
- Line X: <what changed> - <why>
- Line Y-Z: <what changed> - <why>
```cpp
<complete fixed file content>
```

Only output files that need changes. If no issues found, respond with: NO_ISSUES_FOUND'''

SYSTEM_PROMPT_BUILD_ERROR = '''You are analyzing a C++ build error in an IncludeCPP project.

Context:
- IncludeCPP generates pybind11 bindings from .cp plugin files
- SOURCE() and HEADER() link to C++ files
- Code must be in namespace includecpp
- Modules compile to .pyd/.so Python extensions
- Build uses CMake with pybind11

Analyze the error and provide:
1. Root cause (one line)
2. Exact fix with code snippet
3. File and line to modify
4. Prevention tip (one line)

Be concise. No lengthy explanations.'''

SYSTEM_PROMPT_AGENT = '''You are a C++ expert working on an IncludeCPP project.

Context:
- IncludeCPP generates pybind11 bindings
- Code must be in namespace includecpp
- Focus on the specific task requested

Rules:
1. Never remove existing functions unless explicitly asked
2. Preserve public API signatures
3. No comments, docstrings, or explanatory markers
4. Output complete file contents

If removal needed, prefix with:
CONFIRM_REQUIRED: <description>

Output format:
FILE: <path>
```cpp
<complete file content>
```'''


class AIManager:
    def __init__(self):
        self.secret_dir = Path.home() / '.includecpp'
        self.secret_path = self.secret_dir / '.secret'
        self.config = self._load_config()

    def _load_config(self) -> dict:
        if self.secret_path.exists():
            try:
                data = json.loads(self.secret_path.read_text(encoding='utf-8'))
                if 'usage' not in data:
                    data['usage'] = {'total_tokens': 0, 'total_requests': 0, 'last_request': None}
                return data
            except (json.JSONDecodeError, IOError):
                pass
        return {
            'api_key': None,
            'enabled': False,
            'model': DEFAULT_MODEL,
            'usage': {
                'total_tokens': 0,
                'total_requests': 0,
                'last_request': None
            }
        }

    def _save_config(self):
        self.secret_dir.mkdir(parents=True, exist_ok=True)
        self.secret_path.write_text(json.dumps(self.config, indent=2), encoding='utf-8')
        if os.name != 'nt':
            os.chmod(self.secret_path, stat.S_IRUSR | stat.S_IWUSR)

    def _mask_key(self) -> str:
        key = self.config.get('api_key')
        if not key:
            return 'Not set'
        if len(key) <= 8:
            return '****'
        return f"{key[:7]}...{key[-4:]}"

    def set_key(self, key: str) -> Tuple[bool, str]:
        if not key:
            return False, 'API key cannot be empty'
        if not key.startswith('sk-'):
            return False, 'Invalid API key format. Key should start with sk-'
        test_result, test_msg = self._test_key(key)
        if not test_result:
            return False, test_msg
        self.config['api_key'] = key
        self._save_config()
        return True, 'API key saved and verified'

    def _test_key(self, key: str) -> Tuple[bool, str]:
        try:
            headers = {
                'Authorization': f'Bearer {key}',
                'Content-Type': 'application/json'
            }
            data = {
                'model': 'gpt-3.5-turbo',
                'messages': [{'role': 'user', 'content': 'test'}],
                'max_tokens': 1
            }
            response = requests.post(OPENAI_API_URL, headers=headers, json=data, timeout=10)
            if response.status_code == 200:
                return True, 'OK'
            elif response.status_code == 401:
                return False, 'Invalid API key'
            elif response.status_code == 429:
                return True, 'OK (rate limited but valid)'
            else:
                return False, f'API error: {response.status_code}'
        except requests.exceptions.Timeout:
            return False, 'Connection timeout'
        except requests.exceptions.ConnectionError:
            return False, 'Connection failed'
        except Exception as e:
            return False, str(e)

    def is_enabled(self) -> bool:
        return bool(self.config.get('enabled', False) and self.config.get('api_key'))

    def has_key(self) -> bool:
        return bool(self.config.get('api_key'))

    def enable(self) -> Tuple[bool, str]:
        if not self.config.get('api_key'):
            return False, 'No API key configured. Use: includecpp ai key <YOUR_KEY>'
        self.config['enabled'] = True
        self._save_config()
        return True, 'AI features enabled'

    def disable(self) -> Tuple[bool, str]:
        self.config['enabled'] = False
        self._save_config()
        return True, 'AI features disabled'

    def get_model(self) -> str:
        return self.config.get('model', DEFAULT_MODEL)

    def set_model(self, model: str) -> Tuple[bool, str]:
        if model not in MODELS:
            available = ', '.join(MODELS.keys())
            return False, f'Unknown model. Available: {available}'
        self.config['model'] = model
        self._save_config()
        return True, f'Model set to {model}'

    def list_models(self) -> List[Dict[str, Any]]:
        current = self.get_model()
        result = []
        for name, info in MODELS.items():
            result.append({
                'name': name,
                'context': info['context'],
                'active': name == current
            })
        return result

    def get_info(self) -> Dict[str, Any]:
        usage = self.config.get('usage', {})
        return {
            'key_set': bool(self.config.get('api_key')),
            'key_preview': self._mask_key(),
            'enabled': self.config.get('enabled', False),
            'model': self.config.get('model', DEFAULT_MODEL),
            'total_tokens': usage.get('total_tokens', 0),
            'total_requests': usage.get('total_requests', 0),
            'last_request': usage.get('last_request')
        }

    def _update_usage(self, tokens: int):
        if 'usage' not in self.config:
            self.config['usage'] = {'total_tokens': 0, 'total_requests': 0, 'last_request': None}
        self.config['usage']['total_tokens'] = self.config['usage'].get('total_tokens', 0) + tokens
        self.config['usage']['total_requests'] = self.config['usage'].get('total_requests', 0) + 1
        self.config['usage']['last_request'] = datetime.now().isoformat()
        self._save_config()

    def query(self, system_prompt: str, user_prompt: str, temperature: float = 0.3) -> Tuple[bool, str]:
        if not self.config.get('api_key'):
            return False, 'No API key configured'
        model = self.config.get('model', DEFAULT_MODEL)
        model_info = MODELS.get(model, MODELS[DEFAULT_MODEL])
        headers = {
            'Authorization': f'Bearer {self.config["api_key"]}',
            'Content-Type': 'application/json'
        }
        data = {
            'model': model_info['endpoint'],
            'messages': [
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': user_prompt}
            ],
            'temperature': temperature,
            'max_tokens': min(16000, model_info['context'] // 2)
        }
        try:
            response = requests.post(OPENAI_API_URL, headers=headers, json=data, timeout=120)
            if response.status_code == 200:
                result = response.json()
                content = result['choices'][0]['message']['content']
                tokens = result.get('usage', {}).get('total_tokens', 0)
                self._update_usage(tokens)
                return True, content
            elif response.status_code == 401:
                return False, 'Invalid API key'
            elif response.status_code == 429:
                return False, 'Rate limit exceeded. Please wait and try again'
            elif response.status_code == 503:
                return False, 'OpenAI service unavailable'
            else:
                try:
                    err = response.json().get('error', {}).get('message', '')
                except:
                    err = response.text
                return False, f'API error ({response.status_code}): {err}'
        except requests.exceptions.Timeout:
            return False, 'Request timeout'
        except requests.exceptions.ConnectionError:
            return False, 'Connection failed. Check your internet connection'
        except Exception as e:
            return False, str(e)

    def optimize_code(self, files: Dict[str, str], custom_task: Optional[str] = None) -> Tuple[bool, str, List[Dict]]:
        if not files:
            return False, 'No files provided', []
        file_content = ''
        for path, content in files.items():
            file_content += f'\nFILE: {path}\n```cpp\n{content}\n```\n'
        if custom_task:
            prompt = f'Task: {custom_task}\n\nFiles:\n{file_content}'
            system = SYSTEM_PROMPT_AGENT
        else:
            prompt = f'Optimize the following C++ files for performance, safety, and pybind11 compatibility:\n{file_content}'
            system = SYSTEM_PROMPT_OPTIMIZE
        success, response = self.query(system, prompt)
        if not success:
            return False, response, []
        changes = self._parse_file_changes(response)
        return True, response, changes

    def fix_code(self, files: Dict[str, str]) -> Tuple[bool, str, List[Dict]]:
        if not files:
            return False, 'No files provided', []
        file_content = ''
        for path, content in files.items():
            file_content += f'\nFILE: {path}\n```cpp\n{content}\n```\n'
        prompt = f'Analyze and fix issues in these C++ files:\n{file_content}'
        success, response = self.query(SYSTEM_PROMPT_FIX, prompt)
        if not success:
            return False, response, []
        changes = self._parse_file_changes(response)
        return True, response, changes

    def analyze_build_error(self, error: str, source_files: Dict[str, str]) -> Tuple[bool, str]:
        context = f'Build error:\n{error}\n\n'
        if source_files:
            context += 'Relevant source files:\n'
            for path, content in list(source_files.items())[:3]:
                lines = content.split('\n')
                preview = '\n'.join(lines[:100])
                context += f'\nFILE: {path}\n```cpp\n{preview}\n```\n'
        success, response = self.query(SYSTEM_PROMPT_BUILD_ERROR, context)
        return success, response

    def _parse_file_changes(self, response: str) -> List[Dict]:
        changes = []
        lines = response.split('\n')
        current_file = None
        current_content = []
        current_changes_desc = []
        in_code_block = False
        in_changes_section = False
        confirm_required = []
        for line in lines:
            if line.startswith('CONFIRM_REQUIRED:'):
                confirm_required.append(line[17:].strip())
                continue
            if line.startswith('FILE:'):
                if current_file and current_content:
                    changes.append({
                        'file': current_file,
                        'content': '\n'.join(current_content),
                        'changes_desc': current_changes_desc.copy(),
                        'confirm_required': confirm_required.copy()
                    })
                    confirm_required = []
                current_file = line[5:].strip()
                current_content = []
                current_changes_desc = []
                in_code_block = False
                in_changes_section = False
                continue
            if line.startswith('CHANGES:'):
                in_changes_section = True
                continue
            if in_changes_section and not in_code_block:
                stripped = line.strip()
                if stripped.startswith('- '):
                    current_changes_desc.append(stripped[2:])
                elif stripped.startswith('```'):
                    in_changes_section = False
                    in_code_block = True
                continue
            if line.startswith('```cpp'):
                in_code_block = True
                in_changes_section = False
                continue
            if line.startswith('```') and in_code_block:
                in_code_block = False
                continue
            if in_code_block and current_file:
                current_content.append(line)
        if current_file and current_content:
            changes.append({
                'file': current_file,
                'content': '\n'.join(current_content),
                'changes_desc': current_changes_desc,
                'confirm_required': confirm_required
            })
        return changes


_ai_manager_instance = None


def get_ai_manager() -> AIManager:
    global _ai_manager_instance
    if _ai_manager_instance is None:
        _ai_manager_instance = AIManager()
    return _ai_manager_instance
