"""
Supreme Agent Module for Hichi
Multi-agent pipeline for CP/DSA and DBMS problem solving with human-in-the-loop validation.

RATE LIMITS (Gemini 2.5 Pro):
- TPM: 125,000 tokens per minute (rate limit, not context limit)
- Context Window: 1M+ tokens (plenty of room for summaries)
- Typical usage per problem: 10-50k input tokens + 2-10k output tokens
- Well within both rate and context limits for single problem solving
"""

import click
import os
import textwrap
from typing import Optional, Tuple

# === SUPREME CONFIGURATION ===
# API key is now managed securely through `hichi s-key` command
# Password verification is handled by main CLI (jai_shree_ram)

# Model configuration
SUMMARIZATION_MODEL = "gemini-2.0-flash-exp"
CODE_GENERATION_MODEL = "gemini-2.5-pro"

# Check for Google Generative AI library
try:
    import google.generativeai as genai
    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False

# PDF extraction libraries
try:
    import fitz  # PyMuPDF
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False


# Password verification is now handled by CLI (no longer needed here)


def extract_text_from_file(file_path: str) -> Optional[str]:
    """Extract text from PDF or plain text file."""
    file_path = os.path.expanduser(file_path)
    
    if not os.path.exists(file_path):
        click.echo(f"❌ File not found: {file_path}")
        return None
    
    try:
        # Handle PDF files
        if file_path.lower().endswith('.pdf'):
            if not PDF_AVAILABLE:
                click.echo("❌ PDF support not installed. Run: pip install pymupdf")
                return None
            
            text_parts = []
            with fitz.open(file_path) as doc:
                for page_num, page in enumerate(doc, 1):
                    text = page.get_text()
                    if text.strip():
                        text_parts.append(f"--- Page {page_num} ---\n{text}")
            
            return "\n\n".join(text_parts)
        
        # Handle plain text files
        else:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                return f.read()
    
    except Exception as e:
        click.echo(f"❌ Error reading file: {e}")
        return None


def chunk_text(text: str, max_chars: int = 30000) -> list:
    """Split text into manageable chunks for processing."""
    chunks = []
    for i in range(0, len(text), max_chars):
        chunks.append(text[i:i + max_chars])
    return chunks


def call_gemini_api(api_key: str, prompt: str, model_name: str, system_instruction: Optional[str] = None) -> Optional[str]:
    """Call Google Gemini API with the given prompt."""
    if not GEMINI_AVAILABLE:
        return "❌ Google Generative AI library not installed. Run: pip install google-generativeai"
    
    if not api_key or api_key == "YOUR_GEMINI_API_KEY_HERE":
        return "❌ Supreme API key not configured. Run: hichi s-key"
    
    try:
        genai.configure(api_key=api_key)
        
        # Configure model with system instruction if provided
        generation_config = {
            "temperature": 0.7,
            "top_p": 0.95,
            "top_k": 40,
            "max_output_tokens": 8192,
        }
        
        model = genai.GenerativeModel(
            model_name=model_name,
            generation_config=generation_config,
            system_instruction=system_instruction
        )
        
        response = model.generate_content(prompt)
        return response.text
    
    except Exception as e:
        error_msg = str(e)
        if "API_KEY_INVALID" in error_msg or "401" in error_msg:
            return "❌ Invalid Gemini API key. Please update SUPREME_API_KEY in supreme.py"
        elif "RATE_LIMIT" in error_msg or "429" in error_msg:
            return "⏰ Rate limit exceeded. Please try again in a moment."
        elif "NOT_FOUND" in error_msg or "404" in error_msg:
            return f"❌ Model '{model_name}' not found. Check model availability."
        else:
            return f"❌ API Error: {error_msg}"


def summarize_problem(api_key: str, text: str) -> Optional[str]:
    """
    Stage 1: Use Gemini 2.0 Flash to extract and summarize problem statement,
    constraints, input/output format, and requirements.
    """
    click.echo("📖 Stage 1: Extracting and summarizing problem...")
    
    chunks = chunk_text(text)
    chunk_summaries = []
    
    system_instruction = (
        "You are an expert at analyzing competitive programming and database problems. "
        "Extract key information accurately and concisely. "
        "Keep your summary comprehensive but focused - aim for 5,000-10,000 tokens maximum per chunk. 125000 is the TOKEN LIMT"
    )
    
    for idx, chunk in enumerate(chunks, 1):
        if len(chunks) > 1:
            click.echo(f"   Processing chunk {idx}/{len(chunks)}...")
        
        prompt = f"""Analyze this problem statement and extract:

1. **Problem Type**: Is this a CP/DSA problem or a DBMS/SQL problem?
2. **Problem Summary**: Brief description (2-3 sentences)
3. **Constraints**: List all constraints, limits, and edge cases
4. **Input Format**: Describe expected input
5. **Output Format**: Describe expected output
6. **Special Requirements**: Any special conditions or notes
7. **Example**: If present, include one example

Problem text:
{chunk}

Format your response clearly with markdown headers."""
        
        summary = call_gemini_api(api_key, prompt, SUMMARIZATION_MODEL, system_instruction)
        if summary and not summary.startswith("❌"):
            chunk_summaries.append(summary)
        else:
            click.echo(f"⚠️  Warning: Failed to summarize chunk {idx}")
    
    # If multiple chunks, aggregate them
    if len(chunk_summaries) > 1:
        click.echo("   Aggregating summaries...")
        aggregation_prompt = f"""Consolidate these problem analysis summaries into a single coherent summary:

{chr(10).join(f'--- Summary {i+1} ---{chr(10)}{s}' for i, s in enumerate(chunk_summaries))}

Combine them into one clear, complete problem analysis. Remove redundancy but keep all essential constraints and requirements."""
        
        final_summary = call_gemini_api(api_key, aggregation_prompt, SUMMARIZATION_MODEL, system_instruction)
        return final_summary
    elif chunk_summaries:
        return chunk_summaries[0]
    else:
        return None


def generate_solution(api_key: str, problem_summary: str, feedback: Optional[str] = None) -> Optional[str]:
    """
    Stage 2: Use Gemini 2.5 Pro to generate optimized code solution
    (C++ for CP/DSA or SQL for DBMS) based on the approved summary.
    """
    click.echo("🧠 Stage 2: Generating solution with Gemini 2.5 Pro...")
    
    # Determine problem type from summary
    problem_type = "C++"
    if any(keyword in problem_summary.lower() for keyword in ['database', 'sql', 'dbms', 'query', 'table']):
        problem_type = "SQL"
    
    system_instruction = (
        f"You are an expert {problem_type} programmer specializing in "
        f"{'competitive programming and data structures' if problem_type == 'C++' else 'database design and SQL queries'}. "
        "Write clean, efficient with optimal time and space complexity.Don't ADD alotta Comments"
    )
    
    feedback_section = ""
    if feedback:
        feedback_section = f"\n\n**IMPORTANT USER FEEDBACK:**\n{feedback}\n\nPlease incorporate this feedback into your solution.\n"
    
    if problem_type == "C++":
        prompt = f"""Based on the following problem analysis, write a complete, optimized C++ solution.

{problem_summary}{feedback_section}

Requirements:
1. Include all necessary headers
2. Write clean, readable code with meaningful variable names
3. Add comments explaining the approach and key steps
4. Optimize for time and space complexity
5. Handle all edge cases and constraints
6. Include complexity analysis as comments

Please provide ONLY the complete C++ code, properly formatted."""
    
    else:  # SQL
        prompt = f"""Based on the following problem analysis, write optimized SQL queries/schema.

{problem_summary}{feedback_section}

Requirements:
1. Write clean, readable SQL with proper formatting
2. Use appropriate indexes and constraints
3. Add comments explaining complex queries
4. Optimize query performance
5. Handle edge cases
6. Include schema definitions if needed

Please provide ONLY the SQL code, properly formatted."""
    
    solution = call_gemini_api(api_key, prompt, CODE_GENERATION_MODEL, system_instruction)
    return solution


def save_solution(solution: str, problem_type: str) -> None:
    """Save the generated solution to a file."""
    extension = ".cpp" if "C++" in problem_type else ".sql"
    default_name = f"solution{extension}"
    
    if click.confirm(f"\n💾 Save solution to file?", default=True):
        filename = click.prompt("Filename", default=default_name)
        
        # Ensure correct extension
        if not filename.endswith(extension):
            filename += extension
        
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                f.write(solution)
            click.echo(f"✅ Solution saved to: {filename}")
        except Exception as e:
            click.echo(f"❌ Error saving file: {e}")


def display_summary(summary: str) -> None:
    """Display the problem summary in a formatted way."""
    click.echo("\n" + "=" * 80)
    click.echo("📋 PROBLEM ANALYSIS SUMMARY")
    click.echo("=" * 80 + "\n")
    
    # Wrap long lines for better readability
    for line in summary.split('\n'):
        if line.strip():
            if line.startswith('#'):  # Markdown headers
                click.echo(click.style(line, fg='cyan', bold=True))
            else:
                wrapped = textwrap.fill(line, width=78, subsequent_indent='  ')
                click.echo(wrapped)
        else:
            click.echo()
    
    click.echo("\n" + "=" * 80 + "\n")


def supreme_agent_pipeline(api_key: str):
    """
    Main orchestrator for the Supreme Agent pipeline:
    1. Authentication (handled by CLI)
    2. Get problem file path
    3. Extract text from file
    4. Summarize with Gemini 2.0 Flash (Stage 1)
    5. Human-in-the-loop validation + feedback
    6. Generate code with Gemini 2.5 Pro (Stage 2)
    7. Optionally save solution
    """
    click.echo("🌟 " + click.style("HICHI SUPREME AGENT", fg='yellow', bold=True) + " 🌟")
    click.echo("Multi-stage AI pipeline for CP/DSA and DBMS problem solving\n")
    
    # Step 2: Get file path
    file_path = click.prompt("📁 Enter path to problem file (PDF or TXT)", type=str)
    
    # Step 3: Extract text
    click.echo("\n📄 Reading file...")
    text = extract_text_from_file(file_path)
    
    if not text:
        return
    
    click.echo(f"✅ Extracted {len(text)} characters from file\n")
    
    # Step 4: Summarize problem
    summary = summarize_problem(api_key, text)
    
    if not summary or summary.startswith("❌"):
        click.echo(summary or "❌ Failed to generate summary")
        return
    
    click.echo("✅ Summary generated!\n")
    
    # Step 5: Human-in-the-loop validation
    while True:
        display_summary(summary)
        
        if not click.confirm("✋ Do you approve this summary?", default=True):
            click.echo("❌ Summary rejected. Exiting.")
            return
        
        # Ask for feedback
        click.echo("\n💭 You can provide additional feedback or constraints for the solution.")
        click.echo("   (Press Enter to skip)")
        feedback = click.prompt("Your feedback", default="", show_default=False)
        
        if feedback.strip():
            click.echo(f"\n✅ Feedback recorded: {feedback[:100]}...")
        else:
            feedback = None
        
        # Step 6: Generate solution
        solution = generate_solution(api_key, summary, feedback)
        
        if not solution or solution.startswith("❌"):
            click.echo(solution or "❌ Failed to generate solution")
            return
        
        click.echo("\n" + "=" * 80)
        click.echo("💻 GENERATED SOLUTION")
        click.echo("=" * 80 + "\n")
        click.echo(solution)
        click.echo("\n" + "=" * 80 + "\n")
        
        # Ask if user is satisfied
        satisfied = click.confirm("🎯 Are you satisfied with this solution?", default=True)
        
        if satisfied:
            # Determine problem type for saving
            problem_type = "C++" if "```cpp" in solution or "#include" in solution else "SQL"
            save_solution(solution, problem_type)
            click.echo("\n✨ Supreme Agent pipeline completed successfully!")
            break
        else:
            # Allow user to provide more feedback and regenerate
            click.echo("\n🔄 Let's try again with more specific feedback...")
            continue
