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.

IMPORTANT: Be conservative. Only suggest safe optimizations that will NOT break compilation.

Rules:
1. NEVER remove existing functions or code blocks
2. Preserve all public API signatures EXACTLY
3. NEVER change function parameters or return types
4. Only make changes that are 100% safe and backwards compatible
5. Maintain namespace includecpp structure
6. Do not add comments, docstrings, or explanatory text
7. Do not add AI-typical output markers or annotations
8. If unsure about a change, DO NOT make it

Safe optimizations:
- Adding const where appropriate
- Using reserve() before loops
- Replacing raw loops with STL algorithms (if equivalent)
- Adding noexcept to functions that don't throw
- Using move semantics for local variables

AVOID:
- Changing function signatures
- Removing code that might be used elsewhere
- Complex refactoring
- Template changes

If a change would alter existing functionality, you MUST prefix 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 safe optimizations found, 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 inside namespace includecpp { ... }
- Modules compile to .pyd (Windows) or .so (Linux/Mac) Python extensions

CRITICAL CHECKS (in order):
1. SYNTAX ERRORS: Missing semicolons, unmatched braces, invalid tokens
2. NAMESPACE: All classes/functions MUST be inside namespace includecpp { }
3. TYPE ERRORS: Undefined types, wrong return types, parameter mismatches
4. MEMORY: Dangling pointers, leaks, uninitialized variables
5. PYBIND11: Return policies, holder types, opaque types

Rules:
1. Never remove functions or change public API signatures
2. Fix ALL syntax errors first - code must compile
3. Ensure code is inside namespace includecpp { }
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.

IncludeCPP Architecture:
- Plugin files (.cp) define bindings: SOURCE(file.cpp), HEADER(file.h), CLASS(), METHOD(), FUNC()
- SOURCE() links C++ implementation files
- HEADER() links header files for declarations
- All exposed code MUST be in namespace includecpp { }
- Build generates pybind11 bindings in a temp bindings.cpp file
- Compiles to .pyd (Windows) or .so (Linux/Mac) Python extensions

Common Error Categories:
1. USER CODE ERROR: Syntax errors, missing includes, type mismatches in user C++ files
2. PLUGIN FILE ERROR: Wrong signatures in .cp file, missing METHOD(), wrong parameter types
3. NAMESPACE ERROR: Code not inside namespace includecpp { }
4. BINDING ERROR: pybind11 type conversion issues, missing opaque declarations
5. INCLUDECPP BUG: Framework generated incorrect bindings (rare, report at github.com/liliassg/IncludeCPP/issues)

Analyze the error and provide:

ERROR TYPE: [USER CODE | PLUGIN FILE | NAMESPACE | BINDING | INCLUDECPP BUG]

ROOT CAUSE:
<one line explanation>

FIX:
File: <exact file path>
Line: <line number if known>
```cpp
<exact code fix>
```

WHY:
<brief explanation of what was wrong>

PREVENTION:
<one line tip to avoid this in future>'''

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()
        self._doc_cache = None

    def _get_documentation(self) -> str:
        if self._doc_cache:
            return self._doc_cache
        try:
            readme_path = Path(__file__).parent.parent.parent / 'README.md'
            if readme_path.exists():
                self._doc_cache = readme_path.read_text(encoding='utf-8')
                return self._doc_cache
        except:
            pass
        return ''

    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'
        }
        token_limit = min(16000, model_info['context'] // 2)
        data = {
            'model': model_info['endpoint'],
            'messages': [
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': user_prompt}
            ],
            'temperature': temperature
        }
        if model.startswith('gpt-5'):
            data['max_completion_tokens'] = token_limit
        else:
            data['max_tokens'] = token_limit
        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 _build_prompt_with_docs(self, base_prompt: str) -> str:
        docs = self._get_documentation()
        if docs:
            doc_section = docs[:8000] if len(docs) > 8000 else docs
            return f'IncludeCPP Documentation:\n{doc_section}\n\n---\n\n{base_prompt}'
        return base_prompt

    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
        prompt = self._build_prompt_with_docs(prompt)
        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}'
        prompt = self._build_prompt_with_docs(prompt)
        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'
        context = self._build_prompt_with_docs(context)
        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
