# rag/report.py
import json
from typing import Dict, List
from openai import OpenAI
from rag.qa import embed_query
from rag import config

client = OpenAI(api_key=config.OPENAI_API_KEY)


def build_report_from_query(store, query: str, k: int = 100) -> Dict:
    """
    Build a comprehensive 20-25 page report.
    Increased retrieval and detailed section generation.
    """
    qvec = embed_query(query)
    hits = store.search(qvec, k=k)  # Retrieve more chunks

    # Group chunks by document
    doc_chunks = {}
    for hit in hits:
        meta = hit["metadata"]
        doc_id = meta.get("doc_id", "Unknown")
        doc_name = meta.get("doc_name", doc_id)
        
        if doc_id not in doc_chunks:
            doc_chunks[doc_id] = {
                "doc_name": doc_name,
                "chunks": []
            }
        
        doc_chunks[doc_id]["chunks"].append({
            "page": meta.get("page_number", "?"),
            "text": meta.get("text", "")
        })

    # Generate comprehensive sections
    sections = generate_comprehensive_report_sections(doc_chunks, query)
    
    return {
        "query": query,
        "docs": [info["doc_name"] for info in doc_chunks.values()],
        "sections": sections
    }


def generate_comprehensive_report_sections(doc_chunks: Dict, query: str) -> Dict[str, str]:
    """
    Generate detailed, comprehensive report sections for 20-25 pages.
    Each section gets more tokens and detailed instructions.
    """
    sections = {}
    
    # Prepare comprehensive context
    full_context = ""
    for doc_id, info in doc_chunks.items():
        full_context += f"\n\n{'='*60}\n"
        full_context += f"DOCUMENT: {info['doc_name']}\n"
        full_context += f"{'='*60}\n\n"
        
        for chunk in info["chunks"][:30]:  # Use more chunks per doc
            full_context += f"[Page {chunk['page']}]\n{chunk['text']}\n\n"
    
    # Define comprehensive section prompts with more detail requirements
    section_prompts = {
        "1. Executive Summary": {
            "instruction": """Provide a comprehensive executive summary (2-3 pages) that includes:
- Overview of all companies analyzed
- Key highlights and major findings
- Critical risks and opportunities
- Overall financial health assessment
- Strategic recommendations
Be detailed and professional. Include specific numbers, percentages, and facts from the documents.""",
            "max_tokens": 1500
        },
        
        "2. Company Overview & Background": {
            "instruction": """Provide detailed company overview including:
- Company history and establishment
- Business model and operations
- Market position and competitors
- Geographic presence and facilities
- Key milestones and achievements
- Corporate structure
Provide 2-3 paragraphs per company with specific details.""",
            "max_tokens": 1500
        },
        
        "3. Directors & Management": {
            "instruction": """Provide comprehensive analysis of leadership:
- Board of Directors: names, designations, qualifications
- Key Management Personnel with roles
- Management changes during the year
- Experience and expertise of leadership team
- Board composition (independent vs executive directors)
- Board committees and their roles
Present in detailed narrative form with specific names and positions.""",
            "max_tokens": 1500
        },
        
        "4. Shareholding Pattern & Ownership": {
            "instruction": """Detailed shareholding analysis including:
- Promoter shareholding with exact percentages
- Institutional investors and their stakes
- Public shareholding breakdown
- Changes in shareholding during the year
- Top 10 shareholders if available
- Any pledge or encumbrance of shares
Present specific numbers and analyze trends.""",
            "max_tokens": 1500
        },
        
        "5. Financial Performance Analysis": {
            "instruction": """Comprehensive financial analysis covering:
- Revenue trends over 3-5 years with specific numbers
- Profitability metrics (PAT, EBITDA margins)
- Cash flow analysis
- Balance sheet highlights (assets, liabilities, equity)
- Key financial ratios
- Year-over-year growth rates
- Comparison with previous years
Provide detailed numerical analysis with interpretations.""",
            "max_tokens": 1800
        },
        
        "6. Business Operations & Strategy": {
            "instruction": """Detailed operational analysis:
- Core business activities and products/services
- Production capacity and utilization
- Market segments served
- Sales and distribution channels
- Technology and R&D initiatives
- Strategic partnerships and collaborations
- Expansion plans and capital expenditure
Provide comprehensive details with specific examples.""",
            "max_tokens": 1500
        },
        
        "7. Risk Assessment & Management": {
            "instruction": """Comprehensive risk analysis including:
- Business and operational risks
- Financial risks (credit, liquidity, market)
- Regulatory and compliance risks
- Industry-specific risks
- Mitigation strategies for each risk
- Risk management framework
Provide detailed explanation of each risk category.""",
            "max_tokens": 1800
        },
        
        "8. Corporate Governance": {
            "instruction": """Detailed governance analysis:
- Corporate governance framework
- Board meetings and attendance
- Audit committee and other committees
- Internal controls and compliance
- Related party transactions
- Code of conduct policies
- Stakeholder engagement
Provide comprehensive details on governance practices.""",
            "max_tokens": 1500
        },
        
        "9. Regulatory Compliance & Legal": {
            "instruction": """Detailed compliance and legal analysis:
- Regulatory framework applicable
- Compliance status with key regulations
- Legal proceedings if any
- Statutory auditor observations
- Secretarial audit findings
- Penalties or fines if any
- Compliance certifications
Provide specific details on compliance matters.""",
            "max_tokens": 1500
        },
        
        "10. Future Outlook & Recommendations": {
            "instruction": """Comprehensive forward-looking analysis:
- Industry trends and outlook
- Company's strategic initiatives
- Growth opportunities and challenges
- Management's future plans
- Analyst recommendations
- Potential catalysts for growth
- Key monitoring points for stakeholders
Provide detailed forward-looking insights.""",
            "max_tokens": 1500
        }
    }
    
    # Generate each section with LLM
    for section_name, config_dict in section_prompts.items():
        print(f"  Generating: {section_name}...")
        
        prompt = f"""You are a professional financial analyst writing a detailed annual report section.

Report Query: {query}

Context from documents (use ALL relevant information):
{full_context[:15000]}

Section: {section_name}

Task: {config_dict['instruction']}

IMPORTANT:
- Be comprehensive and detailed
- Use specific facts, numbers, and names from the documents
- Cite sources as [Document Name, Page X]
- Write in professional, formal language
- Aim for 2-3 pages of content per section
- Include all relevant details found in the context
- Do not write "based on the context" or similar meta-commentary
- Write as if you are directly presenting the findings

Output:"""
        
        try:
            response = client.chat.completions.create(
                model=config.LLM_MODEL,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=config_dict["max_tokens"]
            )
            
            sections[section_name] = response.choices[0].message.content.strip()
        
        except Exception as e:
            print(f"  ⚠️ Error generating {section_name}: {e}")
            sections[section_name] = f"Error generating this section: {str(e)}"
    
    return sections