post_processor.py

"""
Post-Processor: Split LLM-generated HCL into complete Terraform module
=====================================================================
Takes raw HCL from LLM → splits into main.tf, variables.tf, outputs.tf, README.md

Features:
- Extract hardcoded values → variables.tf with defaults
- Auto-detect outputs from resource blocks
- Generate README with usage instructions
- terraform fmt compatible output
"""

import re
from typing import Optional


def extract_hcl_code(raw_output: str) -> str:
    """Extract HCL code from LLM output (strip markdown fences, explanations)."""
    # Try to find code in ```hcl or ``` blocks
    patterns = [
        r'```(?:hcl|terraform)\s*\n(.*?)```',
        r'```\s*\n(.*?)```',
    ]
    for pat in patterns:
        matches = re.findall(pat, raw_output, re.DOTALL)
        if matches:
            return '\n\n'.join(matches).strip()
    
    # If no fences, try to find resource/data/locals blocks
    lines = raw_output.strip().split('\n')
    code_lines = []
    in_code = False
    brace_depth = 0
    
    for line in lines:
        stripped = line.strip()
        # Start of terraform block
        if re.match(r'^(resource|data|locals|variable|output|module|provider|terraform)\s+', stripped):
            in_code = True
        
        if in_code:
            code_lines.append(line)
            brace_depth += line.count('{') - line.count('}')
            if brace_depth <= 0 and len(code_lines) > 1:
                in_code = False
                brace_depth = 0
                code_lines.append('')
        elif stripped == '' and code_lines:
            pass  # skip blank lines between explanations
    
    result = '\n'.join(code_lines).strip()
    return result if result else raw_output.strip()


def parse_resources(hcl_code: str) -> list[dict]:
    """Parse resource blocks from HCL code."""
    resources = []
    # Match resource "type" "name" { ... }
    pattern = r'resource\s+"([^"]+)"\s+"([^"]+)"'
    for match in re.finditer(pattern, hcl_code):
        resources.append({
            'type': match.group(1),
            'name': match.group(2),
            'full_ref': f'{match.group(1)}.{match.group(2)}'
        })
    return resources


def generate_variables(hcl_code: str, resources: list[dict]) -> str:
    """Generate variables.tf from detected patterns in main.tf."""
    variables = []
    
    # Common variable patterns to extract
    var_patterns = [
        {
            'pattern': r'region\s*=\s*"([\w-]+)"',
            'name': 'aws_region',
            'description': 'AWS region for resources',
            'type': 'string',
        },
        {
            'pattern': r'cidr_block\s*=\s*"([\d./]+)"',
            'name': 'vpc_cidr',
            'description': 'CIDR block for the VPC',
            'type': 'string',
            'match_first': True,
        },
        {
            'pattern': r'instance_type\s*=\s*"([^"]+)"',
            'name': 'instance_type',
            'description': 'EC2 instance type',
            'type': 'string',
        },
        {
            'pattern': r'engine_version\s*=\s*"([^"]+)"',
            'name': 'engine_version',
            'description': 'Database engine version',
            'type': 'string',
        },
        {
            'pattern': r'allocated_storage\s*=\s*(\d+)',
            'name': 'allocated_storage',
            'description': 'Allocated storage in GB',
            'type': 'number',
        },
    ]
    
    seen = set()
    for vp in var_patterns:
        matches = re.findall(vp['pattern'], hcl_code)
        if matches and vp['name'] not in seen:
            default_val = matches[0]
            seen.add(vp['name'])
            
            if vp['type'] == 'number':
                default_line = f'  default     = {default_val}'
            else:
                default_line = f'  default     = "{default_val}"'
            
            variables.append(f'''variable "{vp['name']}" {{
  description = "{vp['description']}"
  type        = {vp['type']}
{default_line}
}}''')
    
    # Always add common variables
    if 'environment' not in seen:
        variables.insert(0, '''variable "environment" {
  description = "Environment name (e.g., dev, staging, prod)"
  type        = string
  default     = "dev"
}''')
    
    if 'project_name' not in seen:
        variables.insert(0, '''variable "project_name" {
  description = "Project name used for resource naming and tagging"
  type        = string
  default     = "my-project"
}''')
    
    if not variables:
        variables.append('# No additional variables detected')
    
    header = '''# ==============================================
# Variables
# ==============================================
# Customize these variables for your environment
# ==============================================
'''
    return header + '\n\n'.join(variables) + '\n'


def generate_outputs(resources: list[dict], hcl_code: str) -> str:
    """Generate outputs.tf from detected resources."""
    outputs = []
    
    # Common output patterns per resource type
    output_map = {
        'aws_vpc': [('vpc_id', 'id', 'The ID of the VPC')],
        'aws_subnet': [('subnet_ids', 'id', 'The IDs of the subnets')],
        'aws_instance': [('instance_id', 'id', 'The ID of the EC2 instance'), ('instance_public_ip', 'public_ip', 'Public IP of the instance')],
        'aws_s3_bucket': [('bucket_name', 'id', 'The name of the S3 bucket'), ('bucket_arn', 'arn', 'The ARN of the S3 bucket')],
        'aws_db_instance': [('db_endpoint', 'endpoint', 'The RDS instance endpoint'), ('db_name', 'db_name', 'The database name')],
        'aws_lambda_function': [('lambda_arn', 'arn', 'The ARN of the Lambda function'), ('lambda_function_name', 'function_name', 'The Lambda function name')],
        'aws_lb': [('lb_dns_name', 'dns_name', 'The DNS name of the load balancer'), ('lb_arn', 'arn', 'The ARN of the load balancer')],
        'aws_ecs_cluster': [('ecs_cluster_name', 'name', 'The ECS cluster name'), ('ecs_cluster_arn', 'arn', 'The ECS cluster ARN')],
        'aws_eks_cluster': [('eks_cluster_endpoint', 'endpoint', 'The EKS cluster endpoint'), ('eks_cluster_name', 'name', 'The EKS cluster name')],
        'aws_api_gateway_rest_api': [('api_gateway_id', 'id', 'The API Gateway REST API ID')],
        'aws_dynamodb_table': [('dynamodb_table_name', 'name', 'The DynamoDB table name'), ('dynamodb_table_arn', 'arn', 'The DynamoDB table ARN')],
        'aws_cloudfront_distribution': [('cloudfront_domain_name', 'domain_name', 'The CloudFront distribution domain name')],
        'aws_sqs_queue': [('sqs_queue_url', 'url', 'The SQS queue URL'), ('sqs_queue_arn', 'arn', 'The SQS queue ARN')],
        'aws_sns_topic': [('sns_topic_arn', 'arn', 'The SNS topic ARN')],
        'aws_iam_role': [('iam_role_arn', 'arn', 'The IAM role ARN')],
        'aws_security_group': [('security_group_id', 'id', 'The security group ID')],
        'aws_elasticache_replication_group': [('redis_endpoint', 'primary_endpoint_address', 'The Redis primary endpoint')],
        'aws_efs_file_system': [('efs_id', 'id', 'The EFS file system ID')],
        'aws_kinesis_stream': [('kinesis_stream_arn', 'arn', 'The Kinesis stream ARN')],
        'aws_kms_key': [('kms_key_arn', 'arn', 'The KMS key ARN'), ('kms_key_id', 'key_id', 'The KMS key ID')],
        'aws_route53_zone': [('route53_zone_id', 'zone_id', 'The Route 53 hosted zone ID')],
        'aws_codepipeline': [('pipeline_arn', 'arn', 'The CodePipeline ARN')],
        'aws_sfn_state_machine': [('state_machine_arn', 'arn', 'The Step Functions state machine ARN')],
        'aws_ecr_repository': [('ecr_repository_url', 'repository_url', 'The ECR repository URL')],
        'aws_cloudtrail': [('cloudtrail_arn', 'arn', 'The CloudTrail ARN')],
        'aws_autoscaling_group': [('asg_name', 'name', 'The Auto Scaling Group name')],
    }
    
    seen_outputs = set()
    for res in resources:
        if res['type'] in output_map:
            for out_name, attr, desc in output_map[res['type']]:
                if out_name not in seen_outputs:
                    seen_outputs.add(out_name)
                    outputs.append(f'''output "{out_name}" {{
  description = "{desc}"
  value       = {res['full_ref']}.{attr}
}}''')
    
    if not outputs:
        # Fallback: output first resource's id
        if resources:
            r = resources[0]
            outputs.append(f'''output "{r['type'].replace('aws_', '')}_id" {{
  description = "The ID of {r['full_ref']}"
  value       = {r['full_ref']}.id
}}''')
    
    header = '''# ==============================================
# Outputs
# ==============================================
# Useful values exported from this module
# ==============================================
'''
    return header + '\n\n'.join(outputs) + '\n'


def generate_readme(prompt: str, resources: list[dict], blueprint_name: str = None) -> str:
    """Generate README.md for the Terraform module."""
    title = blueprint_name or "Terraform Module"
    
    resource_list = '\n'.join(f'- `{r["type"]}` — `{r["name"]}`' for r in resources)
    
    return f'''# {title}

## Description

{prompt}

## Resources Created

{resource_list}

## Usage

```hcl
module "this" {{
  source = "./module"

  project_name = "my-project"
  environment  = "dev"
}}
```

## Requirements

| Name | Version |
|------|---------|
| terraform | >= 1.0 |
| aws | ~> 5.0 |

## Quick Start

```bash
# Initialize Terraform
terraform init

# Preview changes
terraform plan

# Apply infrastructure
terraform apply

# Destroy when done
terraform destroy
```

## Generated by

XAI-Enhanced Terraform Code Generator — Explainable AI techniques 
for LLM-based Infrastructure as Code generation.
'''


def process_output(raw_output: str, original_prompt: str, 
                   blueprint_name: str = None) -> dict:
    """
    Main entry point: process raw LLM output into complete Terraform module.
    
    Returns dict with keys: main_tf, variables_tf, outputs_tf, readme_md, resources
    """
    # Step 1: Extract clean HCL
    main_tf = extract_hcl_code(raw_output)
    
    # Step 2: Remove provider/terraform blocks (pipeline already instructs this)
    main_tf = re.sub(
        r'(?:provider|terraform)\s+"?\w+"?\s*\{[^}]*(?:\{[^}]*\}[^}]*)*\}\s*\n?', 
        '', main_tf
    ).strip()
    
    # Step 3: Parse resources
    resources = parse_resources(main_tf)
    
    # Step 4: Generate supporting files
    variables_tf = generate_variables(main_tf, resources)
    outputs_tf = generate_outputs(resources, main_tf)
    readme_md = generate_readme(original_prompt, resources, blueprint_name)
    
    # Add header comment to main.tf
    main_tf_header = '''# ==============================================
# Main Terraform Configuration
# ==============================================
# Generated by XAI-Enhanced Terraform Generator
# ==============================================

'''
    main_tf = main_tf_header + main_tf + '\n'
    
    return {
        'main_tf': main_tf,
        'variables_tf': variables_tf,
        'outputs_tf': outputs_tf,
        'readme_md': readme_md,
        'resources': [r['full_ref'] for r in resources],
        'resource_count': len(resources),
    }


if __name__ == '__main__':
    # Test
    test_hcl = '''
resource "aws_vpc" "main" {
  cidr_block           = "10.0.0.0/16"
  enable_dns_hostnames = true

  tags = {
    Name = "main-vpc"
  }
}

resource "aws_subnet" "public" {
  vpc_id            = aws_vpc.main.id
  cidr_block        = "10.0.1.0/24"
  availability_zone = "us-east-1a"

  tags = {
    Name = "public-subnet"
  }
}

resource "aws_internet_gateway" "main" {
  vpc_id = aws_vpc.main.id
}
'''
    result = process_output(test_hcl, "Create a VPC with public subnet", "VPC Module")
    for key in ['main_tf', 'variables_tf', 'outputs_tf', 'readme_md']:
        print(f"\n{'='*50}")
        print(f"  {key}")
        print(f"{'='*50}")
        print(result[key][:500])
    print(f"\nResources: {result['resources']}")