from dataclasses import dataclass
from typing import Dict, Optional, List
from enum import Enum

class QueryType(Enum):
    FACTUAL = 'factual'
    EXPLORATORY = 'exploratory'
    COMPARATIVE = 'comparative'
    CODE = 'code'

class DetailLevel(Enum):
    BRIEF = 'brief'
    STANDARD = 'standard'
    DETAILED = 'detailed'
    COMPREHENSIVE = 'comprehensive'

@dataclass
class ReportTemplate:
    template: str
    detail_level: DetailLevel
    query_type: QueryType
    model: Optional[str] = None
    required_sections: Optional[List[str]] = None

    def validate(self) -> bool:
        """Validate that the template contains all required sections"""
        if not self.required_sections:
            return True
        return all(section in self.template for section in self.required_sections)

class ReportTemplateManager:
    def __init__(self):
        self.templates: Dict[str, ReportTemplate] = {}

    def add_template(self, template: ReportTemplate):
        if not template.validate():
            raise ValueError(f"Template missing required sections: {template.required_sections}")
        key = f"{template.query_type.value}_{template.detail_level.value}"
        self.templates[key] = template

    def get_template(self, query_type: QueryType, detail_level: DetailLevel) -> ReportTemplate:
        key = f"{query_type.value}_{detail_level.value}"
        return self.templates.get(key)

    def get_available_templates(self) -> List[str]:
        return list(self.templates.keys())

    def initialize_default_templates(self):
        # Brief templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Summary\n{summary}\n\n## Key Findings\n{key_findings}",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{summary}', '{key_findings}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Research Questions\n{research_questions}\n\n## Key Findings\n{key_findings}",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{research_questions}', '{key_findings}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Key Findings\n{key_findings}",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{comparison_criteria}', '{key_findings}']
        ))
        
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Problem Statement\n{problem_statement}\n\n## Solution\n{solution}\n\n```{language}\n{code_snippet}\n```",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.CODE,
            required_sections=['{title}', '{problem_statement}', '{solution}', '{language}', '{code_snippet}']
        ))

        # Standard templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Introduction\n{introduction}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{introduction}', '{key_findings}', '{analysis}', '{conclusion}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Research Questions\n{research_questions}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{research_questions}', '{methodology}', '{key_findings}', '{analysis}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{comparison_criteria}', '{methodology}', '{key_findings}', '{analysis}']
        ))
        
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Problem Statement\n{problem_statement}\n\n## Approach\n{approach}\n\n## Solution\n{solution}\n\n```{language}\n{code_snippet}\n```\n\n## Explanation\n{explanation}\n\n## Usage Example\n{usage_example}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.CODE,
            required_sections=['{title}', '{problem_statement}', '{approach}', '{solution}', '{language}', '{code_snippet}', '{explanation}', '{usage_example}']
        ))

        # Detailed templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Introduction\n{introduction}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{introduction}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Research Questions\n{research_questions}\n\n## Literature Review\n{literature_review}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{research_questions}', '{literature_review}', '{methodology}', '{key_findings}', '{analysis}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{comparison_criteria}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}']
        ))
        
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Problem Statement\n{problem_statement}\n\n## Context and Requirements\n{context}\n\n## Approach\n{approach}\n\n## Solution\n{solution}\n\n```{language}\n{code_snippet}\n```\n\n## Explanation\n{explanation}\n\n## Alternative Approaches\n{alternatives}\n\n## Best Practices\n{best_practices}\n\n## Usage Examples\n{usage_examples}\n\n## Common Issues\n{common_issues}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.CODE,
            required_sections=['{title}', '{problem_statement}', '{context}', '{approach}', '{solution}', '{language}', '{code_snippet}', '{explanation}', '{alternatives}', '{best_practices}', '{usage_examples}', '{common_issues}']
        ))

        # Comprehensive templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Introduction\n{introduction}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{exec_summary}', '{introduction}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}', '{references}', '{appendices}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Research Questions\n{research_questions}\n\n## Literature Review\n{literature_review}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{exec_summary}', '{research_questions}', '{literature_review}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}', '{references}', '{appendices}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{exec_summary}', '{comparison_criteria}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}', '{references}', '{appendices}']
        ))
        
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Problem Statement\n{problem_statement}\n\n## Technical Background\n{technical_background}\n\n## Architectural Considerations\n{architecture}\n\n## Detailed Solution\n{detailed_solution}\n\n### Implementation Details\n```{language}\n{code_snippet}\n```\n\n## Explanation of Algorithm/Approach\n{algorithm_explanation}\n\n## Performance Considerations\n{performance}\n\n## Alternative Implementations\n{alternatives}\n\n## Best Practices and Design Patterns\n{best_practices}\n\n## Testing and Validation\n{testing}\n\n## Usage Examples\n{usage_examples}\n\n## Common Pitfalls and Workarounds\n{pitfalls}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.CODE,
            required_sections=['{title}', '{exec_summary}', '{problem_statement}', '{technical_background}', '{architecture}', '{detailed_solution}', '{language}', '{code_snippet}', '{algorithm_explanation}', '{performance}', '{alternatives}', '{best_practices}', '{testing}', '{usage_examples}', '{pitfalls}', '{references}', '{appendices}']
        ))
