package generator

import (
	"fmt"
	"strings"

	"github.com/entity-maker/entity-maker/internal/naming"
)

// GenerateModel generates the dataclass model
func GenerateModel(ctx *Context) (string, error) {
	var b strings.Builder

	// Imports
	b.WriteString("from dataclasses import dataclass\n")

	// Check what we need to import
	needsDatetime := NeedsDatetimeImport(ctx.TableInfo.Columns)
	needsDecimal := NeedsDecimalImport(ctx.TableInfo.Columns)

	if needsDatetime {
		b.WriteString("from datetime import datetime\n")
	}
	if needsDecimal {
		b.WriteString("from decimal import Decimal\n")
	}

	b.WriteString("\n")
	b.WriteString("from televend_core.databases.base_model import Base\n")

	// Import related models for foreign keys
	fkImports := make(map[string]string) // module_name -> entity_name
	for _, fk := range ctx.TableInfo.ForeignKeys {
		moduleName := GetRelationshipModuleName(fk.ForeignTableName)
		entityName := GetRelationshipEntityName(fk.ForeignTableName)
		fkImports[moduleName] = entityName
	}

	// Import enum types
	if len(ctx.TableInfo.EnumTypes) > 0 {
		b.WriteString(fmt.Sprintf("from televend_core.databases.televend_repositories.%s.enum import (\n",
			ctx.ModuleName))
		for _, enumType := range ctx.TableInfo.EnumTypes {
			enumName := naming.ToPascalCase(enumType.TypeName)
			b.WriteString(fmt.Sprintf("    %s,\n", enumName))
		}
		b.WriteString(")\n")
	}

	// Write foreign key imports
	for moduleName, entityName := range fkImports {
		b.WriteString(fmt.Sprintf("from televend_core.databases.televend_repositories.%s.model import %s\n",
			moduleName, entityName))
	}

	b.WriteString("\n\n")

	// Class definition
	b.WriteString("@dataclass\n")
	b.WriteString(fmt.Sprintf("class %s(Base):\n", ctx.EntityName))

	// Get required and optional columns
	requiredCols := GetRequiredColumns(ctx.TableInfo.Columns)
	optionalCols := GetOptionalColumns(ctx.TableInfo.Columns)

	// Required fields (non-nullable, no default, not auto-increment, not PK)
	for _, col := range requiredCols {
		fieldName := col.Name
		pythonType := GetPythonTypeForColumn(col, ctx)

		// Regular field
		b.WriteString(fmt.Sprintf("    %s: %s\n", fieldName, pythonType))

		// Add relationship field if this is a foreign key and column ends with _id
		// (to avoid name clashes with FK columns that don't follow _id convention)
		if fk := GetForeignKeyForColumn(col.Name, ctx.TableInfo.ForeignKeys); fk != nil {
			if strings.HasSuffix(col.Name, "_id") {
				relationName := GetRelationshipName(col.Name)
				entityName := GetRelationshipEntityName(fk.ForeignTableName)
				b.WriteString(fmt.Sprintf("    %s: %s\n", relationName, entityName))
			}
		}
	}

	// Empty line between required and optional
	if len(requiredCols) > 0 && len(optionalCols) > 0 {
		b.WriteString("\n")
	}

	// Optional fields
	for _, col := range optionalCols {
		// Skip primary key, we'll add it at the end
		if col.IsPrimaryKey {
			continue
		}

		fieldName := col.Name
		pythonType := GetPythonTypeForColumn(col, ctx)

		b.WriteString(fmt.Sprintf("    %s: %s | None = None\n", fieldName, pythonType))

		// Add relationship field if this is a foreign key and column ends with _id
		// (to avoid name clashes with FK columns that don't follow _id convention)
		if fk := GetForeignKeyForColumn(col.Name, ctx.TableInfo.ForeignKeys); fk != nil {
			if strings.HasSuffix(col.Name, "_id") {
				relationName := GetRelationshipName(col.Name)
				entityName := GetRelationshipEntityName(fk.ForeignTableName)
				b.WriteString(fmt.Sprintf("    %s: %s | None = None\n", relationName, entityName))
			}
		}
	}

	// Add primary key at the end
	for _, col := range ctx.TableInfo.Columns {
		if col.IsPrimaryKey {
			b.WriteString("\n")
			pythonType := GetPythonTypeForColumn(col, ctx)
			b.WriteString(fmt.Sprintf("    %s: %s | None = None\n", col.Name, pythonType))
			break
		}
	}

	return b.String(), nil
}