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>
```'''

SYSTEM_PROMPT_ASK = '''You are an expert assistant for IncludeCPP projects.

You have access to:
- IncludeCPP documentation (README)
- Project source files (.cpp, .h, .py)
- Plugin definitions (.cp files)
- Build system knowledge

Answer questions precisely and concisely. Reference specific files and line numbers when relevant.

Response format:
- Direct answer first
- File references: path:line if applicable
- Code examples if helpful (brief)

Do not explain IncludeCPP basics unless asked. Assume user knows the system.
Max 200 words unless complexity requires more.'''

SYSTEM_PROMPT_EDIT = '''You are a C++ expert editing code in an IncludeCPP project.

Context:
- IncludeCPP generates pybind11 bindings from .cp plugin files
- All exposed code must be in namespace includecpp { }
- Focus on the specific edit task requested

Rules:
1. Make ONLY the requested changes
2. Preserve all existing functionality
3. No comments, docstrings, or explanatory text
4. Maintain exact coding style
5. Do not refactor unrelated code

If the edit would break existing functionality:
CONFIRM_REQUIRED: <description>

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

Only output files that need changes.'''


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 ask_question(self, question: str, files: Dict[str, str], plugins: Dict[str, str] = None) -> Tuple[bool, str]:
        context_parts = []
        if files:
            context_parts.append('Project files:')
            for path, content in files.items():
                lines = content.split('\n')
                preview = '\n'.join(lines[:150])
                ext = Path(path).suffix
                lang = 'cpp' if ext in ['.cpp', '.h', '.hpp', '.c'] else 'python' if ext == '.py' else ''
                context_parts.append(f'\nFILE: {path}\n```{lang}\n{preview}\n```')
        if plugins:
            context_parts.append('\nPlugin definitions:')
            for path, content in plugins.items():
                context_parts.append(f'\nPLUGIN: {path}\n```\n{content}\n```')
        context = '\n'.join(context_parts)
        prompt = f'Question: {question}\n\n{context}'
        prompt = self._build_prompt_with_docs(prompt)
        return self.query(SYSTEM_PROMPT_ASK, prompt, temperature=0.3)

    def edit_code(self, task: str, files: Dict[str, str], think_twice: bool = False) -> Tuple[bool, str, List[Dict]]:
        if not files:
            return False, 'No files provided', []
        file_content = ''
        for path, content in files.items():
            ext = Path(path).suffix
            lang = 'cpp' if ext in ['.cpp', '.h', '.hpp', '.c'] else 'python' if ext == '.py' else ''
            file_content += f'\nFILE: {path}\n```{lang}\n{content}\n```\n'
        prompt = f'Edit task: {task}\n\nFiles:\n{file_content}'
        if think_twice:
            prompt += '\n\nIMPORTANT: Think carefully before making changes. Consider edge cases and potential issues.'
        prompt = self._build_prompt_with_docs(prompt)
        success, response = self.query(SYSTEM_PROMPT_EDIT, prompt)
        if not success:
            return False, response, []
        changes = self._parse_file_changes(response)
        return True, response, changes

    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
